Repository: exo-explore/exo Branch: main Commit: 07598a3af1f6 Files: 724 Total size: 3.1 MB Directory structure: gitextract_vey6weao/ ├── .clauderules ├── .cursorrules ├── .envrc ├── .githooks/ │ ├── post-checkout │ ├── post-commit │ ├── post-merge │ └── pre-push ├── .github/ │ ├── CODEOWNERS │ ├── ISSUE_TEMPLATE/ │ │ ├── bug_report.md │ │ └── feature_request.md │ ├── actions/ │ │ ├── conditional-commit/ │ │ │ └── action.yml │ │ ├── format/ │ │ │ └── action.yml │ │ ├── lint/ │ │ │ └── action.yml │ │ ├── lint-check/ │ │ │ └── action.yml │ │ ├── regenerate-protobufs/ │ │ │ └── action.yml │ │ ├── setup-python-uv/ │ │ │ └── action.yml │ │ ├── unit-test/ │ │ │ └── action.yml │ │ └── verify-clean/ │ │ └── action.yml │ ├── pull_request_template.md │ └── workflows/ │ ├── build-app.yml │ └── pipeline.yml ├── .gitignore ├── .mlx_typings/ │ ├── .gitkeep │ ├── mflux/ │ │ ├── __init__.pyi │ │ ├── callbacks/ │ │ │ ├── __init__.pyi │ │ │ ├── callback.pyi │ │ │ ├── callback_registry.pyi │ │ │ └── generation_context.pyi │ │ ├── cli/ │ │ │ ├── __init__.pyi │ │ │ └── defaults/ │ │ │ └── defaults.pyi │ │ ├── models/ │ │ │ ├── __init__.pyi │ │ │ ├── common/ │ │ │ │ ├── __init__.pyi │ │ │ │ ├── cli/ │ │ │ │ │ └── __init__.pyi │ │ │ │ ├── config/ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ ├── config.pyi │ │ │ │ │ └── model_config.pyi │ │ │ │ ├── latent_creator/ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ └── latent_creator.pyi │ │ │ │ ├── lora/ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ ├── layer/ │ │ │ │ │ │ ├── fused_linear_lora_layer.pyi │ │ │ │ │ │ └── linear_lora_layer.pyi │ │ │ │ │ └── mapping/ │ │ │ │ │ ├── lora_loader.pyi │ │ │ │ │ ├── lora_mapping.pyi │ │ │ │ │ ├── lora_saver.pyi │ │ │ │ │ └── lora_transforms.pyi │ │ │ │ ├── resolution/ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ ├── actions.pyi │ │ │ │ │ ├── config_resolution.pyi │ │ │ │ │ ├── lora_resolution.pyi │ │ │ │ │ ├── path_resolution.pyi │ │ │ │ │ └── quantization_resolution.pyi │ │ │ │ ├── schedulers/ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ ├── base_scheduler.pyi │ │ │ │ │ ├── flow_match_euler_discrete_scheduler.pyi │ │ │ │ │ ├── linear_scheduler.pyi │ │ │ │ │ └── seedvr2_euler_scheduler.pyi │ │ │ │ ├── tokenizer/ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ ├── tokenizer.pyi │ │ │ │ │ ├── tokenizer_loader.pyi │ │ │ │ │ └── tokenizer_output.pyi │ │ │ │ ├── vae/ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ ├── tiling_config.pyi │ │ │ │ │ ├── vae_tiler.pyi │ │ │ │ │ └── vae_util.pyi │ │ │ │ └── weights/ │ │ │ │ ├── __init__.pyi │ │ │ │ ├── loading/ │ │ │ │ │ ├── loaded_weights.pyi │ │ │ │ │ ├── weight_applier.pyi │ │ │ │ │ ├── weight_definition.pyi │ │ │ │ │ └── weight_loader.pyi │ │ │ │ ├── mapping/ │ │ │ │ │ ├── weight_mapper.pyi │ │ │ │ │ ├── weight_mapping.pyi │ │ │ │ │ └── weight_transforms.pyi │ │ │ │ └── saving/ │ │ │ │ └── model_saver.pyi │ │ │ ├── depth_pro/ │ │ │ │ ├── depth_pro_initializer.pyi │ │ │ │ ├── model/ │ │ │ │ │ ├── decoder/ │ │ │ │ │ │ ├── feature_fusion_block_2d.pyi │ │ │ │ │ │ ├── multires_conv_decoder.pyi │ │ │ │ │ │ └── residual_block.pyi │ │ │ │ │ ├── depth_pro.pyi │ │ │ │ │ ├── depth_pro_model.pyi │ │ │ │ │ ├── depth_pro_util.pyi │ │ │ │ │ ├── dino_v2/ │ │ │ │ │ │ ├── attention.pyi │ │ │ │ │ │ ├── dino_vision_transformer.pyi │ │ │ │ │ │ ├── layer_scale.pyi │ │ │ │ │ │ ├── mlp.pyi │ │ │ │ │ │ ├── patch_embed.pyi │ │ │ │ │ │ └── transformer_block.pyi │ │ │ │ │ ├── encoder/ │ │ │ │ │ │ ├── depth_pro_encoder.pyi │ │ │ │ │ │ └── upsample_block.pyi │ │ │ │ │ └── head/ │ │ │ │ │ └── fov_head.pyi │ │ │ │ └── weights/ │ │ │ │ ├── depth_pro_weight_definition.pyi │ │ │ │ └── depth_pro_weight_mapping.pyi │ │ │ ├── fibo/ │ │ │ │ ├── latent_creator/ │ │ │ │ │ └── fibo_latent_creator.pyi │ │ │ │ └── weights/ │ │ │ │ ├── fibo_weight_definition.pyi │ │ │ │ └── fibo_weight_mapping.pyi │ │ │ ├── fibo_vlm/ │ │ │ │ ├── tokenizer/ │ │ │ │ │ ├── qwen2vl_image_processor.pyi │ │ │ │ │ └── qwen2vl_processor.pyi │ │ │ │ └── weights/ │ │ │ │ ├── fibo_vlm_weight_definition.pyi │ │ │ │ └── fibo_vlm_weight_mapping.pyi │ │ │ ├── flux/ │ │ │ │ ├── __init__.pyi │ │ │ │ ├── cli/ │ │ │ │ │ └── __init__.pyi │ │ │ │ ├── flux_initializer.pyi │ │ │ │ ├── latent_creator/ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ └── flux_latent_creator.pyi │ │ │ │ ├── model/ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ ├── flux_text_encoder/ │ │ │ │ │ │ ├── clip_encoder/ │ │ │ │ │ │ │ ├── clip_embeddings.pyi │ │ │ │ │ │ │ ├── clip_encoder.pyi │ │ │ │ │ │ │ ├── clip_encoder_layer.pyi │ │ │ │ │ │ │ ├── clip_mlp.pyi │ │ │ │ │ │ │ ├── clip_sdpa_attention.pyi │ │ │ │ │ │ │ ├── clip_text_model.pyi │ │ │ │ │ │ │ └── encoder_clip.pyi │ │ │ │ │ │ ├── prompt_encoder.pyi │ │ │ │ │ │ └── t5_encoder/ │ │ │ │ │ │ ├── t5_attention.pyi │ │ │ │ │ │ ├── t5_block.pyi │ │ │ │ │ │ ├── t5_dense_relu_dense.pyi │ │ │ │ │ │ ├── t5_encoder.pyi │ │ │ │ │ │ ├── t5_feed_forward.pyi │ │ │ │ │ │ ├── t5_layer_norm.pyi │ │ │ │ │ │ └── t5_self_attention.pyi │ │ │ │ │ ├── flux_transformer/ │ │ │ │ │ │ ├── ada_layer_norm_continuous.pyi │ │ │ │ │ │ ├── ada_layer_norm_zero.pyi │ │ │ │ │ │ ├── ada_layer_norm_zero_single.pyi │ │ │ │ │ │ ├── common/ │ │ │ │ │ │ │ └── attention_utils.pyi │ │ │ │ │ │ ├── embed_nd.pyi │ │ │ │ │ │ ├── feed_forward.pyi │ │ │ │ │ │ ├── guidance_embedder.pyi │ │ │ │ │ │ ├── joint_attention.pyi │ │ │ │ │ │ ├── joint_transformer_block.pyi │ │ │ │ │ │ ├── single_block_attention.pyi │ │ │ │ │ │ ├── single_transformer_block.pyi │ │ │ │ │ │ ├── text_embedder.pyi │ │ │ │ │ │ ├── time_text_embed.pyi │ │ │ │ │ │ ├── timestep_embedder.pyi │ │ │ │ │ │ └── transformer.pyi │ │ │ │ │ ├── flux_vae/ │ │ │ │ │ │ ├── common/ │ │ │ │ │ │ │ ├── attention.pyi │ │ │ │ │ │ │ ├── resnet_block_2d.pyi │ │ │ │ │ │ │ └── unet_mid_block.pyi │ │ │ │ │ │ ├── decoder/ │ │ │ │ │ │ │ ├── conv_in.pyi │ │ │ │ │ │ │ ├── conv_norm_out.pyi │ │ │ │ │ │ │ ├── conv_out.pyi │ │ │ │ │ │ │ ├── decoder.pyi │ │ │ │ │ │ │ ├── up_block_1_or_2.pyi │ │ │ │ │ │ │ ├── up_block_3.pyi │ │ │ │ │ │ │ ├── up_block_4.pyi │ │ │ │ │ │ │ └── up_sampler.pyi │ │ │ │ │ │ ├── encoder/ │ │ │ │ │ │ │ ├── conv_in.pyi │ │ │ │ │ │ │ ├── conv_norm_out.pyi │ │ │ │ │ │ │ ├── conv_out.pyi │ │ │ │ │ │ │ ├── down_block_1.pyi │ │ │ │ │ │ │ ├── down_block_2.pyi │ │ │ │ │ │ │ ├── down_block_3.pyi │ │ │ │ │ │ │ ├── down_block_4.pyi │ │ │ │ │ │ │ ├── down_sampler.pyi │ │ │ │ │ │ │ └── encoder.pyi │ │ │ │ │ │ └── vae.pyi │ │ │ │ │ ├── redux_encoder/ │ │ │ │ │ │ └── redux_encoder.pyi │ │ │ │ │ └── siglip_vision_transformer/ │ │ │ │ │ ├── siglip_encoder.pyi │ │ │ │ │ ├── siglip_encoder_layer.pyi │ │ │ │ │ ├── siglip_mlp.pyi │ │ │ │ │ ├── siglip_multi_head_attention_pooling_head.pyi │ │ │ │ │ ├── siglip_sdpa_attention.pyi │ │ │ │ │ ├── siglip_vision_embeddings.pyi │ │ │ │ │ └── siglip_vision_transformer.pyi │ │ │ │ ├── variants/ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ ├── concept_attention/ │ │ │ │ │ │ ├── attention_data.pyi │ │ │ │ │ │ ├── joint_attention_concept.pyi │ │ │ │ │ │ ├── joint_transformer_block_concept.pyi │ │ │ │ │ │ └── transformer_concept.pyi │ │ │ │ │ ├── controlnet/ │ │ │ │ │ │ └── transformer_controlnet.pyi │ │ │ │ │ ├── kontext/ │ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ │ ├── flux_kontext.pyi │ │ │ │ │ │ └── kontext_util.pyi │ │ │ │ │ └── txt2img/ │ │ │ │ │ └── flux.pyi │ │ │ │ └── weights/ │ │ │ │ ├── __init__.pyi │ │ │ │ ├── flux_lora_mapping.pyi │ │ │ │ ├── flux_weight_definition.pyi │ │ │ │ └── flux_weight_mapping.pyi │ │ │ ├── qwen/ │ │ │ │ ├── __init__.pyi │ │ │ │ ├── cli/ │ │ │ │ │ └── __init__.pyi │ │ │ │ ├── latent_creator/ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ └── qwen_latent_creator.pyi │ │ │ │ ├── model/ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ ├── qwen_text_encoder/ │ │ │ │ │ │ ├── qwen_attention.pyi │ │ │ │ │ │ ├── qwen_encoder.pyi │ │ │ │ │ │ ├── qwen_encoder_layer.pyi │ │ │ │ │ │ ├── qwen_mlp.pyi │ │ │ │ │ │ ├── qwen_patch_merger.pyi │ │ │ │ │ │ ├── qwen_prompt_encoder.pyi │ │ │ │ │ │ ├── qwen_rms_norm.pyi │ │ │ │ │ │ ├── qwen_rope.pyi │ │ │ │ │ │ ├── qwen_text_encoder.pyi │ │ │ │ │ │ ├── qwen_vision_attention.pyi │ │ │ │ │ │ ├── qwen_vision_block.pyi │ │ │ │ │ │ ├── qwen_vision_language_encoder.pyi │ │ │ │ │ │ ├── qwen_vision_mlp.pyi │ │ │ │ │ │ ├── qwen_vision_patch_embed.pyi │ │ │ │ │ │ ├── qwen_vision_rotary_embedding.pyi │ │ │ │ │ │ └── qwen_vision_transformer.pyi │ │ │ │ │ ├── qwen_transformer/ │ │ │ │ │ │ ├── qwen_attention.pyi │ │ │ │ │ │ ├── qwen_feed_forward.pyi │ │ │ │ │ │ ├── qwen_rope.pyi │ │ │ │ │ │ ├── qwen_time_text_embed.pyi │ │ │ │ │ │ ├── qwen_timestep_embedding.pyi │ │ │ │ │ │ ├── qwen_timesteps.pyi │ │ │ │ │ │ ├── qwen_transformer.pyi │ │ │ │ │ │ ├── qwen_transformer_block.pyi │ │ │ │ │ │ └── qwen_transformer_rms_norm.pyi │ │ │ │ │ └── qwen_vae/ │ │ │ │ │ ├── qwen_image_attention_block_3d.pyi │ │ │ │ │ ├── qwen_image_causal_conv_3d.pyi │ │ │ │ │ ├── qwen_image_decoder_3d.pyi │ │ │ │ │ ├── qwen_image_down_block_3d.pyi │ │ │ │ │ ├── qwen_image_encoder_3d.pyi │ │ │ │ │ ├── qwen_image_mid_block_3d.pyi │ │ │ │ │ ├── qwen_image_res_block_3d.pyi │ │ │ │ │ ├── qwen_image_resample_3d.pyi │ │ │ │ │ ├── qwen_image_rms_norm.pyi │ │ │ │ │ ├── qwen_image_up_block_3d.pyi │ │ │ │ │ └── qwen_vae.pyi │ │ │ │ ├── qwen_initializer.pyi │ │ │ │ ├── tokenizer/ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ ├── qwen_image_processor.pyi │ │ │ │ │ ├── qwen_vision_language_processor.pyi │ │ │ │ │ └── qwen_vision_language_tokenizer.pyi │ │ │ │ ├── variants/ │ │ │ │ │ ├── __init__.pyi │ │ │ │ │ ├── edit/ │ │ │ │ │ │ ├── qwen_edit_util.pyi │ │ │ │ │ │ └── qwen_image_edit.pyi │ │ │ │ │ └── txt2img/ │ │ │ │ │ └── qwen_image.pyi │ │ │ │ └── weights/ │ │ │ │ ├── __init__.pyi │ │ │ │ ├── qwen_lora_mapping.pyi │ │ │ │ ├── qwen_weight_definition.pyi │ │ │ │ └── qwen_weight_mapping.pyi │ │ │ ├── seedvr2/ │ │ │ │ └── weights/ │ │ │ │ ├── seedvr2_weight_definition.pyi │ │ │ │ └── seedvr2_weight_mapping.pyi │ │ │ └── z_image/ │ │ │ ├── latent_creator/ │ │ │ │ └── z_image_latent_creator.pyi │ │ │ └── weights/ │ │ │ ├── z_image_weight_definition.pyi │ │ │ └── z_image_weight_mapping.pyi │ │ ├── release/ │ │ │ └── __init__.pyi │ │ └── utils/ │ │ ├── __init__.pyi │ │ ├── box_values.pyi │ │ ├── exceptions.pyi │ │ ├── generated_image.pyi │ │ ├── image_util.pyi │ │ ├── metadata_builder.pyi │ │ └── version_util.pyi │ ├── mlx/ │ │ ├── core/ │ │ │ ├── __init__.pyi │ │ │ ├── cuda/ │ │ │ │ └── __init__.pyi │ │ │ ├── distributed/ │ │ │ │ └── __init__.pyi │ │ │ ├── metal/ │ │ │ │ └── __init__.pyi │ │ │ └── random/ │ │ │ └── __init__.pyi │ │ ├── nn/ │ │ │ ├── __init__.pyi │ │ │ ├── init.pyi │ │ │ ├── layers/ │ │ │ │ ├── __init__.pyi │ │ │ │ ├── activations.pyi │ │ │ │ ├── base.pyi │ │ │ │ ├── containers.pyi │ │ │ │ ├── convolution.pyi │ │ │ │ ├── convolution_transpose.pyi │ │ │ │ ├── distributed.pyi │ │ │ │ ├── dropout.pyi │ │ │ │ ├── embedding.pyi │ │ │ │ ├── linear.pyi │ │ │ │ ├── normalization.pyi │ │ │ │ ├── pooling.pyi │ │ │ │ ├── positional_encoding.pyi │ │ │ │ ├── quantized.pyi │ │ │ │ ├── recurrent.pyi │ │ │ │ ├── transformer.pyi │ │ │ │ └── upsample.pyi │ │ │ ├── losses.pyi │ │ │ └── utils.pyi │ │ └── utils.pyi │ └── mlx_lm/ │ ├── __init__.pyi │ ├── convert.pyi │ ├── generate.pyi │ ├── models/ │ │ ├── __init__.pyi │ │ ├── base.pyi │ │ ├── bitlinear_layers.pyi │ │ ├── cache.pyi │ │ ├── deepseek_v3.pyi │ │ ├── glm4_moe.pyi │ │ ├── glm_moe_dsa.pyi │ │ ├── nemotron_h.pyi │ │ ├── qwen3_5.pyi │ │ ├── qwen3_5_moe.pyi │ │ ├── qwen3_next.pyi │ │ ├── step3p5.pyi │ │ └── switch_layers.pyi │ ├── sample_utils.pyi │ ├── tokenizer_utils.pyi │ └── utils.pyi ├── .python-version ├── .swift-format ├── .vscode/ │ ├── extensions.json │ └── settings.json ├── .zed/ │ └── settings.json ├── AGENTS.md ├── CONTRIBUTING.md ├── Cargo.toml ├── LICENSE ├── MISSED_THINGS.md ├── PLATFORMS.md ├── README.md ├── RULES.md ├── TODO.md ├── app/ │ └── EXO/ │ ├── EXO/ │ │ ├── Assets.xcassets/ │ │ │ ├── AccentColor.colorset/ │ │ │ │ └── Contents.json │ │ │ ├── AppIcon.appiconset/ │ │ │ │ └── Contents.json │ │ │ ├── Contents.json │ │ │ └── menubar-icon.imageset/ │ │ │ └── Contents.json │ │ ├── ContentView.swift │ │ ├── EXO.entitlements │ │ ├── EXOApp.swift │ │ ├── ExoProcessController.swift │ │ ├── Info.plist │ │ ├── Models/ │ │ │ └── ClusterState.swift │ │ ├── Preview Content/ │ │ │ └── Preview Assets.xcassets/ │ │ │ └── Contents.json │ │ ├── Services/ │ │ │ ├── BugReportService.swift │ │ │ ├── ClusterStateService.swift │ │ │ ├── LocalNetworkChecker.swift │ │ │ ├── NetworkSetupHelper.swift │ │ │ ├── NetworkStatusService.swift │ │ │ ├── ThunderboltBridgeDetector.swift │ │ │ └── ThunderboltBridgeService.swift │ │ ├── ViewModels/ │ │ │ ├── InstanceViewModel.swift │ │ │ └── NodeViewModel.swift │ │ ├── Views/ │ │ │ ├── FirstLaunchPopout.swift │ │ │ ├── InstanceRowView.swift │ │ │ ├── NodeDetailView.swift │ │ │ ├── NodeRowView.swift │ │ │ ├── SettingsView.swift │ │ │ ├── SettingsWindowController.swift │ │ │ └── TopologyMiniView.swift │ │ └── main.swift │ ├── EXO.xcodeproj/ │ │ ├── project.pbxproj │ │ ├── project.xcworkspace/ │ │ │ ├── contents.xcworkspacedata │ │ │ └── xcshareddata/ │ │ │ └── swiftpm/ │ │ │ └── Package.resolved │ │ └── xcshareddata/ │ │ └── xcschemes/ │ │ └── EXO.xcscheme │ ├── EXOTests/ │ │ └── EXOTests.swift │ ├── EXOUITests/ │ │ ├── EXOUITests.swift │ │ └── EXOUITestsLaunchTests.swift │ └── uninstall-exo.sh ├── bench/ │ ├── bench.toml │ ├── eval_configs/ │ │ └── models.toml │ ├── eval_tool_calls.py │ ├── exo_bench.py │ ├── exo_eval.py │ ├── harness.py │ ├── parallel_requests.py │ ├── pyproject.toml │ ├── scenarios.toml │ ├── single-m3-ultra.toml │ ├── src/ │ │ └── exo_bench/ │ │ └── __init__.py │ └── vendor/ │ ├── __init__.py │ └── lcb_testing_util.py ├── dashboard/ │ ├── dashboard.nix │ ├── package.json │ ├── parts.nix │ ├── src/ │ │ ├── app.css │ │ ├── app.d.ts │ │ ├── app.html │ │ ├── lib/ │ │ │ ├── components/ │ │ │ │ ├── ChatAttachments.svelte │ │ │ │ ├── ChatForm.svelte │ │ │ │ ├── ChatMessages.svelte │ │ │ │ ├── ChatModelSelector.svelte │ │ │ │ ├── ChatSidebar.svelte │ │ │ │ ├── ConnectionBanner.svelte │ │ │ │ ├── DeviceIcon.svelte │ │ │ │ ├── FamilyLogos.svelte │ │ │ │ ├── FamilySidebar.svelte │ │ │ │ ├── HeaderNav.svelte │ │ │ │ ├── HuggingFaceResultItem.svelte │ │ │ │ ├── ImageLightbox.svelte │ │ │ │ ├── ImageParamsPanel.svelte │ │ │ │ ├── MarkdownContent.svelte │ │ │ │ ├── ModelCard.svelte │ │ │ │ ├── ModelFilterPopover.svelte │ │ │ │ ├── ModelPickerGroup.svelte │ │ │ │ ├── ModelPickerModal.svelte │ │ │ │ ├── PrefillProgressBar.svelte │ │ │ │ ├── ToastContainer.svelte │ │ │ │ ├── TokenHeatmap.svelte │ │ │ │ ├── TopologyGraph.svelte │ │ │ │ └── index.ts │ │ │ ├── stores/ │ │ │ │ ├── app.svelte.ts │ │ │ │ ├── favorites.svelte.ts │ │ │ │ ├── recents.svelte.ts │ │ │ │ └── toast.svelte.ts │ │ │ ├── types/ │ │ │ │ └── files.ts │ │ │ └── utils/ │ │ │ └── downloads.ts │ │ └── routes/ │ │ ├── +layout.svelte │ │ ├── +page.svelte │ │ ├── downloads/ │ │ │ └── +page.svelte │ │ └── traces/ │ │ ├── +page.svelte │ │ └── [taskId]/ │ │ └── +page.svelte │ ├── svelte.config.js │ ├── tsconfig.json │ └── vite.config.ts ├── docs/ │ ├── api.md │ └── architecture.md ├── flake.nix ├── justfile ├── nix/ │ ├── apple-sdk/ │ │ └── metadata/ │ │ └── versions.json │ ├── apple-sdk-overlay.nix │ ├── darwin-build-fixes.patch │ ├── metal-toolchain.nix │ └── mlx.nix ├── packaging/ │ ├── dmg/ │ │ ├── create-dmg.sh │ │ └── generate-background.py │ └── pyinstaller/ │ └── exo.spec ├── pyproject.toml ├── python/ │ └── parts.nix ├── resources/ │ ├── image_model_cards/ │ │ ├── exolabs--FLUX.1-Kontext-dev-4bit.toml │ │ ├── exolabs--FLUX.1-Kontext-dev-8bit.toml │ │ ├── exolabs--FLUX.1-Kontext-dev.toml │ │ ├── exolabs--FLUX.1-Krea-dev-4bit.toml │ │ ├── exolabs--FLUX.1-Krea-dev-8bit.toml │ │ ├── exolabs--FLUX.1-Krea-dev.toml │ │ ├── exolabs--FLUX.1-dev-4bit.toml │ │ ├── exolabs--FLUX.1-dev-8bit.toml │ │ ├── exolabs--FLUX.1-dev.toml │ │ ├── exolabs--FLUX.1-schnell-4bit.toml │ │ ├── exolabs--FLUX.1-schnell-8bit.toml │ │ ├── exolabs--FLUX.1-schnell.toml │ │ ├── exolabs--Qwen-Image-4bit.toml │ │ ├── exolabs--Qwen-Image-8bit.toml │ │ ├── exolabs--Qwen-Image-Edit-2509-4bit.toml │ │ ├── exolabs--Qwen-Image-Edit-2509-8bit.toml │ │ ├── exolabs--Qwen-Image-Edit-2509.toml │ │ └── exolabs--Qwen-Image.toml │ └── inference_model_cards/ │ ├── mlx-community--DeepSeek-V3.1-4bit.toml │ ├── mlx-community--DeepSeek-V3.1-8bit.toml │ ├── mlx-community--GLM-4.5-Air-8bit.toml │ ├── mlx-community--GLM-4.5-Air-bf16.toml │ ├── mlx-community--GLM-4.7-4bit.toml │ ├── mlx-community--GLM-4.7-6bit.toml │ ├── mlx-community--GLM-4.7-8bit-gs32.toml │ ├── mlx-community--GLM-4.7-Flash-4bit.toml │ ├── mlx-community--GLM-4.7-Flash-5bit.toml │ ├── mlx-community--GLM-4.7-Flash-6bit.toml │ ├── mlx-community--GLM-4.7-Flash-8bit.toml │ ├── mlx-community--GLM-5-8bit.toml │ ├── mlx-community--GLM-5-MXFP4-Q8.toml │ ├── mlx-community--GLM-5-bf16.toml │ ├── mlx-community--Kimi-K2-Instruct-4bit.toml │ ├── mlx-community--Kimi-K2-Thinking.toml │ ├── mlx-community--Kimi-K2.5.toml │ ├── mlx-community--Llama-3.1-Nemotron-70B-Instruct-HF-4bit.toml │ ├── mlx-community--Llama-3.1-Nemotron-70B-Instruct-HF-8bit.toml │ ├── mlx-community--Llama-3.1-Nemotron-70B-Instruct-HF-bf16.toml │ ├── mlx-community--Llama-3.1-Nemotron-Nano-4B-v1.1-4bit.toml │ ├── mlx-community--Llama-3.1-Nemotron-Nano-4B-v1.1-8bit.toml │ ├── mlx-community--Llama-3.1-Nemotron-Nano-4B-v1.1-bf16.toml │ ├── mlx-community--Llama-3.2-1B-Instruct-4bit.toml │ ├── mlx-community--Llama-3.2-3B-Instruct-4bit.toml │ ├── mlx-community--Llama-3.2-3B-Instruct-8bit.toml │ ├── mlx-community--Llama-3.3-70B-Instruct-4bit.toml │ ├── mlx-community--Llama-3.3-70B-Instruct-8bit.toml │ ├── mlx-community--Meta-Llama-3.1-70B-Instruct-4bit.toml │ ├── mlx-community--Meta-Llama-3.1-8B-Instruct-4bit.toml │ ├── mlx-community--Meta-Llama-3.1-8B-Instruct-8bit.toml │ ├── mlx-community--Meta-Llama-3.1-8B-Instruct-bf16.toml │ ├── mlx-community--MiniMax-M2.1-3bit.toml │ ├── mlx-community--MiniMax-M2.1-8bit.toml │ ├── mlx-community--MiniMax-M2.5-4bit.toml │ ├── mlx-community--MiniMax-M2.5-6bit.toml │ ├── mlx-community--MiniMax-M2.5-8bit.toml │ ├── mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-4Bit.toml │ ├── mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-5Bit.toml │ ├── mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-6Bit.toml │ ├── mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-8Bit.toml │ ├── mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-BF16.toml │ ├── mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-MXFP4.toml │ ├── mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4.toml │ ├── mlx-community--NVIDIA-Nemotron-Nano-9B-v2-4bits.toml │ ├── mlx-community--NVIDIA-Nemotron-Nano-9B-v2-6bit.toml │ ├── mlx-community--Qwen3-0.6B-4bit.toml │ ├── mlx-community--Qwen3-0.6B-8bit.toml │ ├── mlx-community--Qwen3-235B-A22B-Instruct-2507-4bit.toml │ ├── mlx-community--Qwen3-235B-A22B-Instruct-2507-8bit.toml │ ├── mlx-community--Qwen3-30B-A3B-4bit.toml │ ├── mlx-community--Qwen3-30B-A3B-8bit.toml │ ├── mlx-community--Qwen3-Coder-480B-A35B-Instruct-4bit.toml │ ├── mlx-community--Qwen3-Coder-480B-A35B-Instruct-8bit.toml │ ├── mlx-community--Qwen3-Coder-Next-4bit.toml │ ├── mlx-community--Qwen3-Coder-Next-5bit.toml │ ├── mlx-community--Qwen3-Coder-Next-6bit.toml │ ├── mlx-community--Qwen3-Coder-Next-8bit.toml │ ├── mlx-community--Qwen3-Coder-Next-bf16.toml │ ├── mlx-community--Qwen3-Next-80B-A3B-Instruct-4bit.toml │ ├── mlx-community--Qwen3-Next-80B-A3B-Instruct-8bit.toml │ ├── mlx-community--Qwen3-Next-80B-A3B-Thinking-4bit.toml │ ├── mlx-community--Qwen3-Next-80B-A3B-Thinking-8bit.toml │ ├── mlx-community--Qwen3.5-122B-A10B-4bit.toml │ ├── mlx-community--Qwen3.5-122B-A10B-6bit.toml │ ├── mlx-community--Qwen3.5-122B-A10B-8bit.toml │ ├── mlx-community--Qwen3.5-122B-A10B-bf16.toml │ ├── mlx-community--Qwen3.5-27B-4bit.toml │ ├── mlx-community--Qwen3.5-27B-8bit.toml │ ├── mlx-community--Qwen3.5-2B-MLX-8bit.toml │ ├── mlx-community--Qwen3.5-35B-A3B-4bit.toml │ ├── mlx-community--Qwen3.5-35B-A3B-8bit.toml │ ├── mlx-community--Qwen3.5-397B-A17B-4bit.toml │ ├── mlx-community--Qwen3.5-397B-A17B-6bit.toml │ ├── mlx-community--Qwen3.5-397B-A17B-8bit.toml │ ├── mlx-community--Qwen3.5-9B-4bit.toml │ ├── mlx-community--Qwen3.5-9B-8bit.toml │ ├── mlx-community--Step-3.5-Flash-4bit.toml │ ├── mlx-community--Step-3.5-Flash-6bit.toml │ ├── mlx-community--Step-3.5-Flash-8Bit.toml │ ├── mlx-community--gpt-oss-120b-MXFP4-Q8.toml │ ├── mlx-community--gpt-oss-20b-MXFP4-Q8.toml │ └── mlx-community--llama-3.3-70b-instruct-fp16.toml ├── rust/ │ ├── exo_pyo3_bindings/ │ │ ├── Cargo.toml │ │ ├── README.md │ │ ├── exo_pyo3_bindings.pyi │ │ ├── pyproject.toml │ │ ├── src/ │ │ │ ├── allow_threading.rs │ │ │ ├── bin/ │ │ │ │ └── stub_gen.rs │ │ │ ├── ident.rs │ │ │ ├── lib.rs │ │ │ └── networking.rs │ │ └── tests/ │ │ ├── dummy.rs │ │ └── test_python.py │ ├── networking/ │ │ ├── Cargo.toml │ │ ├── examples/ │ │ │ └── chatroom.rs │ │ ├── src/ │ │ │ ├── RESEARCH_NOTES.txt │ │ │ ├── discovery.rs │ │ │ ├── lib.rs │ │ │ └── swarm.rs │ │ └── tests/ │ │ └── dummy.rs │ ├── parts.nix │ └── util/ │ ├── Cargo.toml │ └── src/ │ ├── lib.rs │ └── wakerdeque.rs ├── scripts/ │ └── fetch_kv_heads.py ├── src/ │ └── exo/ │ ├── __init__.py │ ├── __main__.py │ ├── api/ │ │ ├── __init__.py │ │ ├── adapters/ │ │ │ ├── __init__.py │ │ │ ├── chat_completions.py │ │ │ ├── claude.py │ │ │ ├── ollama.py │ │ │ └── responses.py │ │ ├── main.py │ │ ├── tests/ │ │ │ ├── test_api_error_handling.py │ │ │ ├── test_cancel_command.py │ │ │ ├── test_claude_api.py │ │ │ ├── test_claude_tool_use.py │ │ │ └── test_openai_responses_api.py │ │ └── types/ │ │ ├── __init__.py │ │ ├── api.py │ │ ├── claude_api.py │ │ ├── ollama_api.py │ │ └── openai_responses.py │ ├── download/ │ │ ├── coordinator.py │ │ ├── download_utils.py │ │ ├── huggingface_utils.py │ │ ├── impl_shard_downloader.py │ │ ├── shard_downloader.py │ │ └── tests/ │ │ ├── __init__.py │ │ ├── test_download_verification.py │ │ ├── test_offline_mode.py │ │ └── test_re_download.py │ ├── main.py │ ├── master/ │ │ ├── __init__.py │ │ ├── image_store.py │ │ ├── main.py │ │ ├── placement.py │ │ ├── placement_utils.py │ │ └── tests/ │ │ ├── conftest.py │ │ ├── test_master.py │ │ ├── test_placement.py │ │ ├── test_placement_utils.py │ │ └── test_topology.py │ ├── routing/ │ │ ├── __init__.py │ │ ├── connection_message.py │ │ ├── event_router.py │ │ ├── router.py │ │ ├── tests/ │ │ │ └── test_event_buffer.py │ │ └── topics.py │ ├── shared/ │ │ ├── __init__.py │ │ ├── apply.py │ │ ├── constants.py │ │ ├── election.py │ │ ├── logging.py │ │ ├── models/ │ │ │ └── model_cards.py │ │ ├── tests/ │ │ │ ├── __init__.py │ │ │ ├── conftest.py │ │ │ ├── test_apply/ │ │ │ │ ├── test_apply_node_download.py │ │ │ │ └── test_apply_runner_deleted.py │ │ │ ├── test_election.py │ │ │ ├── test_node_id_persistence.py │ │ │ ├── test_resolve_reasoning_params.py │ │ │ ├── test_state_serialization.py │ │ │ └── test_xdg_paths.py │ │ ├── topology.py │ │ ├── tracing.py │ │ └── types/ │ │ ├── __init__.py │ │ ├── chunks.py │ │ ├── commands.py │ │ ├── common.py │ │ ├── events.py │ │ ├── memory.py │ │ ├── mlx.py │ │ ├── multiaddr.py │ │ ├── profiling.py │ │ ├── state.py │ │ ├── tasks.py │ │ ├── text_generation.py │ │ ├── thunderbolt.py │ │ ├── topology.py │ │ └── worker/ │ │ ├── downloads.py │ │ ├── instances.py │ │ ├── runner_response.py │ │ ├── runners.py │ │ └── shards.py │ ├── utils/ │ │ ├── __init__.py │ │ ├── banner.py │ │ ├── channels.py │ │ ├── dashboard_path.py │ │ ├── disk_event_log.py │ │ ├── event_buffer.py │ │ ├── fs.py │ │ ├── info_gatherer/ │ │ │ ├── __init__.py │ │ │ ├── info_gatherer.py │ │ │ ├── macmon.py │ │ │ ├── net_profile.py │ │ │ ├── system_info.py │ │ │ └── tests/ │ │ │ └── test_tb_parsing.py │ │ ├── keyed_backoff.py │ │ ├── phantom.py │ │ ├── power_sampler.py │ │ ├── pydantic_ext.py │ │ ├── reactive.py │ │ ├── task_group.py │ │ └── tests/ │ │ ├── test_event_log.py │ │ ├── test_mp_channel.py │ │ ├── test_power_sampler.py │ │ └── test_tagged.py │ └── worker/ │ ├── __init__.py │ ├── engines/ │ │ ├── __init__.py │ │ ├── image/ │ │ │ ├── __init__.py │ │ │ ├── config.py │ │ │ ├── distributed_model.py │ │ │ ├── generate.py │ │ │ ├── models/ │ │ │ │ ├── __init__.py │ │ │ │ ├── base.py │ │ │ │ ├── flux/ │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── adapter.py │ │ │ │ │ ├── config.py │ │ │ │ │ ├── kontext_adapter.py │ │ │ │ │ └── wrappers.py │ │ │ │ └── qwen/ │ │ │ │ ├── __init__.py │ │ │ │ ├── adapter.py │ │ │ │ ├── config.py │ │ │ │ ├── edit_adapter.py │ │ │ │ └── wrappers.py │ │ │ └── pipeline/ │ │ │ ├── __init__.py │ │ │ ├── block_wrapper.py │ │ │ ├── kv_cache.py │ │ │ └── runner.py │ │ └── mlx/ │ │ ├── __init__.py │ │ ├── auto_parallel.py │ │ ├── cache.py │ │ ├── constants.py │ │ ├── dsml_encoding.py │ │ ├── generator/ │ │ │ ├── __init__.py │ │ │ ├── batch_generate.py │ │ │ └── generate.py │ │ ├── tests/ │ │ │ └── test_batch_generate.py │ │ └── utils_mlx.py │ ├── main.py │ ├── plan.py │ ├── runner/ │ │ ├── __init__.py │ │ ├── bootstrap.py │ │ ├── image_models/ │ │ │ ├── __init__.py │ │ │ └── runner.py │ │ ├── llm_inference/ │ │ │ ├── __init__.py │ │ │ ├── batch_generator.py │ │ │ ├── model_output_parsers.py │ │ │ ├── runner.py │ │ │ └── tool_parsers.py │ │ └── runner_supervisor.py │ └── tests/ │ ├── TODO.tests │ ├── __init__.py │ ├── constants.py │ └── unittests/ │ ├── __init__.py │ ├── conftest.py │ ├── test_download/ │ │ └── __init__.py │ ├── test_mlx/ │ │ ├── __init__.py │ │ ├── conftest.py │ │ ├── test_auto_parallel.py │ │ ├── test_batch_vs_generate.py │ │ ├── test_distributed_fix.py │ │ ├── test_kv_prefix_cache.py │ │ ├── test_pipeline_prefill_callbacks.py │ │ ├── test_prefix_cache_architectures.py │ │ └── test_tokenizers.py │ ├── test_plan/ │ │ ├── __init__.py │ │ ├── test_download_and_loading.py │ │ ├── test_runner_lifecycle.py │ │ ├── test_task_forwarding.py │ │ └── test_warmup.py │ └── test_runner/ │ ├── __init__.py │ ├── test_dsml_e2e.py │ ├── test_event_ordering.py │ ├── test_glm_tool_parsing.py │ ├── test_parse_gpt_oss.py │ ├── test_parse_tool_calls.py │ └── test_runner_supervisor.py ├── tests/ │ ├── auto_bench.sh │ ├── eval_tool_calls.sh │ ├── get_all_models_on_cluster.py │ ├── headless_runner.py │ ├── run_exo_on.sh │ └── start_distributed_test.py └── tmp/ ├── config_examples/ │ ├── claude_code.sh │ └── opencode.json ├── gen_card.py ├── prompt.txt ├── quantize_and_upload.py ├── run_llm.py ├── run_llm.sh ├── set_rdma_network_config.sh └── test_trust_remote_code_attack.sh ================================================ FILE CONTENTS ================================================ ================================================ FILE: .clauderules ================================================ # Claude Code Rules - Follow Every Rule Exactly You must prioritize straightforward code semantics, well-named types, clear function signatures, and robust, carefully-chosen abstractions. Think about how your decisions might impact these aspects of code quality before proposing any changes. You have access to all modern Python features from Python 3.13, 3.12, 3.11... **When you're done making changes, remove any redundant comments; remaining comments should only apply to complex code segments, adding relevant context.** ## 1. Code Discipline * Eliminate superfluous `try`/`catch` and `if` branches through strict typing and static analysis. * Use pure functions unless you must mutate fixed state—then wrap that state in a class. * Every function is **referentially transparent**: same inputs ⇒ same outputs, no hidden state, no unintended I/O. * Put side-effects in injectable "effect handlers"; keep core logic pure. ## 2. Naming * Choose descriptive, non-abbreviated names—no 3-letter acronyms or non-standard contractions. * Anyone reading a function's type signature alone should grasp its purpose without extra context. ## 3. Typing * Maintain **strict, exhaustive** typing; never bypass the type-checker. * Default to `Literal[...]` when an enum-like set is needed. * Prefer built-in types; when two values share structure but differ in meaning, enforce separation: * Use `typing.NewType` for primitives (zero runtime cost). * For serializable objects, add a `type: str` field that states the object's identity. ## 4. Pydantic * Read, respect, and rely on Pydantic documentation. * Centralize a common `ConfigDict` with `frozen=True` and `strict=True` (or stricter) and reuse it everywhere. * For hierarchies of `BaseModel` variants, declare a discriminated union with `typing.Annotated[Base, Field(discriminator='variant')]`; publish a single `TypeAdapter[Base]` so all variants share one strict validator. ## 5. IDs & UUIDs * Subclass Pydantic's `UUID4` for custom ID types. * Generate fresh IDs with `uuid.uuid4()`. * Create idempotency keys by hashing *persisted* state plus a **function-specific salt** to avoid collisions after crashes. ## 6. Error Handling * Catch an exception **only** where you can handle or transform it meaningfully. * State in the docstring **where** each exception is expected to be handled and **why**. ## 7. Dependencies * Introduce new external dependencies only after approval. * Request only libraries common in production environments. ## 8. Use of `@final` & Freezing * Mark classes, methods, and variables as `@final` or otherwise immutable wherever applicable. ## 9. Repository Workflow If you spot a rule violation within code that you've not been asked to work on directly, inform the user rather than patching it ad-hoc. --- ### One-Sentence Summary Write strictly-typed, pure, self-describing Python that uses Pydantic, well-scoped side-effects, immutable state, approved dependencies, and explicit error handling. ================================================ FILE: .cursorrules ================================================ # follow **every** rule exactly; report any violation instead of silently fixing it. You must prioritize straightforward code semantics, well-named types, clear function signatures, and robust, carefully-chosen abstractions. Think about how your decisions might impact these aspects of code quality before proposing any changes. You can use the advanced features of `typing`. You have access to all of the new features from Python 3.13, 3.12, 3.11... **When you're done making your changes, remove any redundant comments that you may have left; the comments that remain should only apply to complex segments of code, adding relevant context.** ## 1. Code Discipline * Eliminate superfluous `try` / `catch` and `if` branches through strict typing and static analysis. * Use pure functions unless you must mutate fixed state—then wrap that state in a class. * Every function is **referentially transparent**: same inputs ⇒ same outputs, no hidden state, no unintended I/O. * Put side-effects in injectable “effect handlers”; keep core logic pure. ## 2. Naming * Choose descriptive, non-abbreviated names—no 3-letter acronyms or non-standard contractions. * Anyone reading a function’s type signature alone should grasp its purpose without extra context. ## 3. Typing * Maintain **strict, exhaustive** typing; never bypass the type-checker. * Default to `Literal[...]` when an enum-like set is needed. * Prefer built-in types; when two values share structure but differ in meaning, enforce separation: * Use `typing.NewType` for primitives (zero runtime cost). * For serialisable objects, add a `type: str` field that states the object’s identity. ## 4. Pydantic * Read, respect, and rely on Pydantic docs. * Centralise a common `ConfigDict` with `frozen=True` and `strict=True` (or stricter) and reuse it everywhere. * For hierarchies of `BaseModel` variants, declare a discriminated union with `typing.Annotated[Base, Field(discriminator='variant')]`; publish a single `TypeAdapter[Base]` so all variants share one strict validator. ## 5. IDs & UUIDs * Subclass Pydantic’s `UUID4` for custom ID types. * Generate fresh IDs with `uuid.uuid4()`. * Create idempotency keys by hashing *persisted* state plus a **function-specific salt** to avoid collisions after crashes. ## 6. Error Handling * Catch an exception **only** where you can handle or transform it meaningfully. * State in the docstring **where** each exception is expected to be handled and **why**. ## 7. Dependencies * Introduce new external dependencies only after approval. * Request only libraries common in production environments. ## 8. Use of `@final` & Freezing * Mark classes, methods, and variables as `@final` or otherwise immutable wherever applicable. ## 9. Repository Workflow If you spot a rule violation within code that you've not been asked to work on directly, inform the user rather than patching it ad-hoc. --- ### One-Sentence Summary Write strictly-typed, pure, self-describing Python that uses Pydantic, well-scoped side-effects, immutable state, approved dependencies, and explicit error handling ================================================ FILE: .envrc ================================================ use flake ================================================ FILE: .githooks/post-checkout ================================================ #!/bin/sh command -v git-lfs >/dev/null 2>&1 || { printf >&2 "\n%s\n\n" "This repository is configured for Git LFS but 'git-lfs' was not found on your path. If you no longer wish to use Git LFS, remove this hook by deleting the 'post-checkout' file in the hooks directory (set by 'core.hookspath'; usually '.git/hooks')."; exit 2; } git lfs post-checkout "$@" ================================================ FILE: .githooks/post-commit ================================================ #!/bin/sh command -v git-lfs >/dev/null 2>&1 || { printf >&2 "\n%s\n\n" "This repository is configured for Git LFS but 'git-lfs' was not found on your path. If you no longer wish to use Git LFS, remove this hook by deleting the 'post-commit' file in the hooks directory (set by 'core.hookspath'; usually '.git/hooks')."; exit 2; } git lfs post-commit "$@" ================================================ FILE: .githooks/post-merge ================================================ #!/bin/sh command -v git-lfs >/dev/null 2>&1 || { printf >&2 "\n%s\n\n" "This repository is configured for Git LFS but 'git-lfs' was not found on your path. If you no longer wish to use Git LFS, remove this hook by deleting the 'post-merge' file in the hooks directory (set by 'core.hookspath'; usually '.git/hooks')."; exit 2; } git lfs post-merge "$@" ================================================ FILE: .githooks/pre-push ================================================ #!/bin/sh command -v git-lfs >/dev/null 2>&1 || { printf >&2 "\n%s\n\n" "This repository is configured for Git LFS but 'git-lfs' was not found on your path. If you no longer wish to use Git LFS, remove this hook by deleting the 'pre-push' file in the hooks directory (set by 'core.hookspath'; usually '.git/hooks')."; exit 2; } git lfs pre-push "$@" ================================================ FILE: .github/CODEOWNERS ================================================ * @ToxicPine * @AlexCheema * @GeluVrabie ================================================ FILE: .github/ISSUE_TEMPLATE/bug_report.md ================================================ --- name: Bug Report about: Create a report to help us improve title: '[BUG] ' labels: bug assignees: '' --- ## Describe the bug A clear and concise description of what the bug is. ## To Reproduce Steps to reproduce the behavior: 1. 2. 3. ## Expected behavior A clear and concise description of what you expected to happen. ## Actual behavior A clear and concise description of what actually happened. ## Environment - macOS Version: - EXO Version: - Hardware: - Device 1: (e.g., MacBook Pro M1 Max, 32GB RAM) - Device 2: (e.g., Mac Mini M2, 16GB RAM) - Additional devices: - Interconnection: - (e.g., Thunderbolt 4 cable between Device 1 and 2) - (e.g., WiFi 6 for Device 3) - (e.g., 10GbE Ethernet between all devices) ## Additional context Add any other context about the problem here. ================================================ FILE: .github/ISSUE_TEMPLATE/feature_request.md ================================================ --- name: Feature Request about: Suggest an idea for this project title: '' labels: enhancement assignees: '' --- Describe what you'd like to see added to EXO. ================================================ FILE: .github/actions/conditional-commit/action.yml ================================================ name: Commit if changed description: "Create a commit when the working tree is dirty" inputs: message: description: "Commit message" required: true runs: using: composite steps: - name: Commit changed files shell: bash run: | git diff --quiet && exit 0 git commit -am "${{ inputs.message }}" ================================================ FILE: .github/actions/format/action.yml ================================================ name: Format Code description: "Run code formatter" runs: using: "composite" steps: - name: Format code run: nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just fmt shell: bash ================================================ FILE: .github/actions/lint/action.yml ================================================ name: Lint Code description: "Run code linter" runs: using: "composite" steps: - name: Lint code run: nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just lint shell: bash ================================================ FILE: .github/actions/lint-check/action.yml ================================================ name: Lint Check description: "Check for lint errors" runs: using: "composite" steps: - name: Lint check run: nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just lint-check shell: bash ================================================ FILE: .github/actions/regenerate-protobufs/action.yml ================================================ name: Regenerate Protobufs description: "Regenerate protobuf files" runs: using: "composite" steps: - name: Regenerate protobufs run: nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just regenerate-protobufs shell: bash ================================================ FILE: .github/actions/setup-python-uv/action.yml ================================================ name: Setup Python & uv description: "Regenerate Python environment from uv.lock" runs: using: "composite" steps: - name: Install uv uses: astral-sh/setup-uv@v6 with: enable-cache: true cache-dependency-glob: uv.lock - name: Install Python run: uv python install shell: bash - name: Sync run: uv sync --locked --all-extras --dev shell: bash ================================================ FILE: .github/actions/unit-test/action.yml ================================================ name: Unit Test description: "Run unit tests" runs: using: "composite" steps: - name: Run unit tests run: | nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just sync-clean nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just test-fast shell: bash ================================================ FILE: .github/actions/verify-clean/action.yml ================================================ name: Verify Clean Working Tree description: "Fail the job if the previous step left the working tree dirty" inputs: step: description: "The name of the step that just executed" required: true runs: using: composite steps: - name: Check git diff shell: bash run: | if ! git diff --quiet; then echo "Error: ${{ inputs.step }} left working tree dirty." >&2 git --no-pager diff >&2 exit 1 fi ================================================ FILE: .github/pull_request_template.md ================================================ ## Motivation ## Changes ## Why It Works ## Test Plan ### Manual Testing ### Automated Testing ================================================ FILE: .github/workflows/build-app.yml ================================================ name: Build EXO macOS DMG # Release workflow: # 1. Create a draft GitHub Release with the tag name (e.g. v1.0.0) and write release notes in markdown # 2. Push the tag: git tag v1.0.0 && git push origin v1.0.0 # 3. This workflow builds, signs, and notarizes the DMG # 4. Release notes are embedded in appcast.xml for Sparkle (rendered as markdown) # 5. DMG and appcast.xml are uploaded to S3 # 6. The draft GitHub Release is published with the DMG attached # # For alpha releases (e.g. v1.0.0-alpha.1): draft release and notes are optional. # If no draft exists, a release is auto-created with generated notes. on: workflow_dispatch: push: tags: - "v*" branches: - "test-app" jobs: build-macos-app: runs-on: "macos-26" permissions: contents: write env: SPARKLE_VERSION: 2.9.0-beta.1 SPARKLE_DOWNLOAD_PREFIX: ${{ secrets.SPARKLE_DOWNLOAD_PREFIX }} SPARKLE_FEED_URL: ${{ secrets.SPARKLE_FEED_URL }} SPARKLE_ED25519_PUBLIC: ${{ secrets.SPARKLE_ED25519_PUBLIC }} SPARKLE_ED25519_PRIVATE: ${{ secrets.SPARKLE_ED25519_PRIVATE }} SPARKLE_S3_BUCKET: ${{ secrets.SPARKLE_S3_BUCKET }} SPARKLE_S3_PREFIX: ${{ secrets.SPARKLE_S3_PREFIX }} EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT: ${{ secrets.EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT }} AWS_REGION: ${{ secrets.AWS_REGION }} EXO_BUILD_NUMBER: ${{ github.run_number }} EXO_LIBP2P_NAMESPACE: ${{ github.ref_name }} steps: # ============================================================ # Checkout and tag validation # ============================================================ - name: Checkout uses: actions/checkout@v4 with: fetch-depth: 0 - name: Derive release version from tag run: | if [[ "$GITHUB_REF_NAME" == "test-app" || "${{ github.event_name }}" == "workflow_dispatch" ]]; then VERSION="0.0.0-alpha.0" echo "IS_ALPHA=true" >> $GITHUB_ENV else VERSION="${GITHUB_REF_NAME#v}" if [[ "$VERSION" == *-alpha* ]]; then echo "IS_ALPHA=true" >> $GITHUB_ENV else echo "IS_ALPHA=false" >> $GITHUB_ENV fi fi echo "RELEASE_VERSION=$VERSION" >> $GITHUB_ENV - name: Compute build version from semver run: | VERSION="$RELEASE_VERSION" # Extract major.minor.patch (strip prerelease suffix) BASE_VERSION="${VERSION%%-*}" MAJOR=$(echo "$BASE_VERSION" | cut -d. -f1) MINOR=$(echo "$BASE_VERSION" | cut -d. -f2) PATCH=$(echo "$BASE_VERSION" | cut -d. -f3) # Extract prerelease number (e.g., "alpha.2" -> 2, or 999 for releases) if [[ "$VERSION" == *-* ]]; then PRERELEASE_PART="${VERSION#*-}" PRERELEASE_NUM="${PRERELEASE_PART##*.}" # Default to 0 if not a number if ! [[ "$PRERELEASE_NUM" =~ ^[0-9]+$ ]]; then PRERELEASE_NUM=0 fi else PRERELEASE_NUM=999 fi # Compute: PRERELEASE + (1000 * PATCH) + (1_000_000 * MINOR) + (1_000_000_000 * MAJOR) BUILD_VERSION=$((PRERELEASE_NUM + 1000 * PATCH + 1000000 * MINOR + 1000000000 * MAJOR)) echo "EXO_BUILD_VERSION=$BUILD_VERSION" >> $GITHUB_ENV echo "Computed build version: $BUILD_VERSION from $VERSION" - name: Ensure tag commit is on main if: github.ref_type == 'tag' run: | git fetch origin main # Alpha tags can be on any branch, production tags must be on main if [[ "$IS_ALPHA" == "true" ]]; then echo "Alpha tag detected, skipping main branch check" elif ! git merge-base --is-ancestor origin/main HEAD; then echo "Production tag must point to a commit on main" exit 1 fi - name: Fetch and validate release notes if: github.ref_type == 'tag' env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} run: | # Find draft release by name using gh release list (more reliable with default token) echo "Looking for draft release named '$GITHUB_REF_NAME'..." DRAFT_EXISTS=$(gh release list --json name,isDraft --jq ".[] | select(.isDraft == true) | select(.name == \"$GITHUB_REF_NAME\") | .name" 2>/dev/null || echo "") if [[ -z "$DRAFT_EXISTS" ]]; then if [[ "$IS_ALPHA" == "true" ]]; then echo "No draft release found for alpha tag $GITHUB_REF_NAME (optional for alphas)" echo "HAS_RELEASE_NOTES=false" >> $GITHUB_ENV exit 0 fi echo "ERROR: No draft release found for tag $GITHUB_REF_NAME" echo "Please create a draft release with release notes before pushing the tag." exit 1 fi # Fetch full release details via API to get body and ID echo "Found draft release, fetching details..." RELEASE_JSON=$(gh api repos/${{ github.repository }}/releases --jq ".[] | select(.draft == true) | select(.name == \"$GITHUB_REF_NAME\")" 2>/dev/null || echo "") # Extract release notes NOTES=$(echo "$RELEASE_JSON" | jq -r '.body // ""') if [[ -z "$NOTES" || "$NOTES" == "null" ]]; then if [[ "$IS_ALPHA" == "true" ]]; then echo "Draft release has no notes (optional for alphas)" echo "HAS_RELEASE_NOTES=false" >> $GITHUB_ENV exit 0 fi echo "ERROR: Draft release exists but has no release notes" echo "Please add release notes to the draft release before pushing the tag." exit 1 fi # Save release ID for later publishing RELEASE_ID=$(echo "$RELEASE_JSON" | jq -r '.id') echo "DRAFT_RELEASE_ID=$RELEASE_ID" >> $GITHUB_ENV echo "HAS_RELEASE_NOTES=true" >> $GITHUB_ENV echo "Found draft release (ID: $RELEASE_ID), saving release notes..." echo "$NOTES" > /tmp/release_notes.md echo "RELEASE_NOTES_FILE=/tmp/release_notes.md" >> $GITHUB_ENV # ============================================================ # Install dependencies # ============================================================ - name: Select Xcode 26.2 run: | sudo xcode-select -s /Applications/Xcode_26.2.app if ! xcrun -f metal >/dev/null 2>&1; then echo "Metal toolchain is not installed." exit 1 fi - name: Install Homebrew packages run: brew install just awscli macmon - name: Install UV uses: astral-sh/setup-uv@v6 with: enable-cache: true cache-dependency-glob: uv.lock - name: Setup Python run: | uv python install uv sync --locked - name: Install Nix uses: cachix/install-nix-action@v31 with: nix_path: nixpkgs=channel:nixos-unstable - name: Configure Cachix uses: cachix/cachix-action@v14 with: name: exo authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}" - name: Build dashboard run: | DASHBOARD_OUT=$(nix build .#dashboard --print-build-logs --no-link --print-out-paths) mkdir -p dashboard/build cp -r "$DASHBOARD_OUT"/* dashboard/build/ - name: Install Sparkle CLI run: | CLI_URL="${SPARKLE_CLI_URL:-https://github.com/sparkle-project/Sparkle/releases/download/${SPARKLE_VERSION}/Sparkle-${SPARKLE_VERSION}.tar.xz}" echo "Downloading Sparkle CLI from: $CLI_URL" mkdir -p /tmp/sparkle curl --fail --location --output /tmp/sparkle.tar.xz "$CLI_URL" tar -xJf /tmp/sparkle.tar.xz -C /tmp/sparkle --strip-components=1 echo "SPARKLE_BIN=/tmp/sparkle/bin" >> $GITHUB_ENV - name: Prepare code-signing keychain env: MACOS_CERTIFICATE: ${{ secrets.MACOS_CERTIFICATE }} MACOS_CERTIFICATE_PASSWORD: ${{ secrets.MACOS_CERTIFICATE_PASSWORD }} PROVISIONING_PROFILE: ${{ secrets.PROVISIONING_PROFILE }} run: | KEYCHAIN_PATH="$HOME/Library/Keychains/build.keychain-db" # Create fresh keychain security create-keychain -p "$MACOS_CERTIFICATE_PASSWORD" "$KEYCHAIN_PATH" # Disable auto-lock (no timeout, no lock-on-sleep) security set-keychain-settings "$KEYCHAIN_PATH" # Add to search list while preserving existing keychains security list-keychains -d user -s "$KEYCHAIN_PATH" $(security list-keychains -d user | tr -d '"') # Set as default and unlock security default-keychain -s "$KEYCHAIN_PATH" security unlock-keychain -p "$MACOS_CERTIFICATE_PASSWORD" "$KEYCHAIN_PATH" # Import certificate with full access for codesign echo "$MACOS_CERTIFICATE" | base64 --decode > /tmp/cert.p12 security import /tmp/cert.p12 -k "$KEYCHAIN_PATH" -P "$MACOS_CERTIFICATE_PASSWORD" \ -T /usr/bin/codesign -T /usr/bin/security -T /usr/bin/productbuild rm /tmp/cert.p12 # Allow codesign to access the key without prompting security set-key-partition-list -S apple-tool:,apple:,codesign: -s -k "$MACOS_CERTIFICATE_PASSWORD" "$KEYCHAIN_PATH" # Verify keychain is unlocked and identity is available echo "Verifying signing identity..." security find-identity -v -p codesigning "$KEYCHAIN_PATH" # Setup provisioning profile mkdir -p "$HOME/Library/Developer/Xcode/UserData/Provisioning Profiles" echo "$PROVISIONING_PROFILE" | base64 --decode > "$HOME/Library/Developer/Xcode/UserData/Provisioning Profiles/EXO.provisionprofile" # Export keychain path for other steps echo "BUILD_KEYCHAIN_PATH=$KEYCHAIN_PATH" >> $GITHUB_ENV # ============================================================ # Build the bundle # ============================================================ - name: Build PyInstaller bundle run: uv run pyinstaller packaging/pyinstaller/exo.spec - name: Build Swift app env: MACOS_CERTIFICATE_PASSWORD: ${{ secrets.MACOS_CERTIFICATE_PASSWORD }} SPARKLE_FEED_URL: ${{ secrets.SPARKLE_FEED_URL }} SPARKLE_ED25519_PUBLIC: ${{ secrets.SPARKLE_ED25519_PUBLIC }} run: | cd app/EXO security unlock-keychain -p "$MACOS_CERTIFICATE_PASSWORD" "$BUILD_KEYCHAIN_PATH" SIGNING_IDENTITY=$(security find-identity -v -p codesigning "$BUILD_KEYCHAIN_PATH" | awk -F '"' '{print $2}') xcodebuild clean build \ -scheme EXO \ -configuration Release \ -derivedDataPath build \ MARKETING_VERSION="$RELEASE_VERSION" \ CURRENT_PROJECT_VERSION="$EXO_BUILD_VERSION" \ EXO_BUILD_TAG="$RELEASE_VERSION" \ EXO_BUILD_COMMIT="$GITHUB_SHA" \ SPARKLE_FEED_URL="$SPARKLE_FEED_URL" \ SPARKLE_ED25519_PUBLIC="$SPARKLE_ED25519_PUBLIC" \ EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT="$EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT" \ CODE_SIGNING_IDENTITY="$SIGNING_IDENTITY" \ CODE_SIGN_INJECT_BASE_ENTITLEMENTS=YES mkdir -p ../../output cp -R build/Build/Products/Release/EXO.app ../../output/EXO.app - name: Inject PyInstaller runtime run: | rm -rf output/EXO.app/Contents/Resources/exo mkdir -p output/EXO.app/Contents/Resources cp -R dist/exo output/EXO.app/Contents/Resources/exo - name: Codesign PyInstaller runtime env: MACOS_CERTIFICATE_PASSWORD: ${{ secrets.MACOS_CERTIFICATE_PASSWORD }} run: | cd output security unlock-keychain -p "$MACOS_CERTIFICATE_PASSWORD" "$BUILD_KEYCHAIN_PATH" SIGNING_IDENTITY=$(security find-identity -v -p codesigning "$BUILD_KEYCHAIN_PATH" | awk -F '"' '{print $2}') RUNTIME_DIR="EXO.app/Contents/Resources/exo" find "$RUNTIME_DIR" -type f \( -perm -111 -o -name "*.dylib" -o -name "*.so" \) -print0 | while IFS= read -r -d '' file; do /usr/bin/codesign --force --timestamp --options runtime \ --sign "$SIGNING_IDENTITY" "$file" done - name: Sign, notarize, and create DMG env: MACOS_CERTIFICATE_PASSWORD: ${{ secrets.MACOS_CERTIFICATE_PASSWORD }} APPLE_NOTARIZATION_USERNAME: ${{ secrets.APPLE_NOTARIZATION_USERNAME }} APPLE_NOTARIZATION_PASSWORD: ${{ secrets.APPLE_NOTARIZATION_PASSWORD }} APPLE_NOTARIZATION_TEAM: ${{ secrets.APPLE_NOTARIZATION_TEAM }} run: | cd output security unlock-keychain -p "$MACOS_CERTIFICATE_PASSWORD" "$BUILD_KEYCHAIN_PATH" SIGNING_IDENTITY=$(security find-identity -v -p codesigning "$BUILD_KEYCHAIN_PATH" | awk -F '"' '{print $2}') /usr/bin/codesign --deep --force --timestamp --options runtime \ --sign "$SIGNING_IDENTITY" EXO.app mkdir -p dmg-root cp -R EXO.app dmg-root/ ln -s /Applications dmg-root/Applications DMG_NAME="EXO-${RELEASE_VERSION}.dmg" hdiutil create -volname "EXO" -srcfolder dmg-root -ov -format UDZO "$DMG_NAME" /usr/bin/codesign --force --timestamp --options runtime \ --sign "$SIGNING_IDENTITY" "$DMG_NAME" if [[ -n "$APPLE_NOTARIZATION_USERNAME" ]]; then SUBMISSION_OUTPUT=$(xcrun notarytool submit "$DMG_NAME" \ --apple-id "$APPLE_NOTARIZATION_USERNAME" \ --password "$APPLE_NOTARIZATION_PASSWORD" \ --team-id "$APPLE_NOTARIZATION_TEAM" \ --wait --timeout 15m 2>&1) echo "$SUBMISSION_OUTPUT" SUBMISSION_ID=$(echo "$SUBMISSION_OUTPUT" | awk 'tolower($1)=="id:" && $2 ~ /^[0-9a-fA-F-]+$/ {print $2; exit}') STATUS=$(echo "$SUBMISSION_OUTPUT" | awk 'tolower($1)=="status:" {print $2; exit}') if [[ -n "$SUBMISSION_ID" ]]; then xcrun notarytool log "$SUBMISSION_ID" \ --apple-id "$APPLE_NOTARIZATION_USERNAME" \ --password "$APPLE_NOTARIZATION_PASSWORD" \ --team-id "$APPLE_NOTARIZATION_TEAM" > notarization-log.txt || true echo "===== Notarization Log =====" cat notarization-log.txt echo "============================" fi if [[ "$STATUS" != "Accepted" ]]; then echo "Notarization failed with status: ${STATUS:-Unknown}" exit 1 fi xcrun stapler staple "$DMG_NAME" fi - name: Generate Sparkle appcast env: SPARKLE_DOWNLOAD_PREFIX: ${{ env.SPARKLE_DOWNLOAD_PREFIX }} SPARKLE_ED25519_PRIVATE: ${{ secrets.SPARKLE_ED25519_PRIVATE }} IS_ALPHA: ${{ env.IS_ALPHA }} run: | set -euo pipefail cd output DOWNLOAD_PREFIX="${SPARKLE_DOWNLOAD_PREFIX:-https://assets.exolabs.net}" echo "$SPARKLE_ED25519_PRIVATE" > sparkle_ed25519.key chmod 600 sparkle_ed25519.key CHANNEL_FLAG="" if [[ "$IS_ALPHA" == "true" ]]; then CHANNEL_FLAG="--channel alpha" echo "Generating appcast for alpha channel" fi $SPARKLE_BIN/generate_appcast \ --ed-key-file sparkle_ed25519.key \ --download-url-prefix "$DOWNLOAD_PREFIX" \ $CHANNEL_FLAG \ . - name: Inject release notes into appcast if: github.ref_type == 'tag' && env.HAS_RELEASE_NOTES == 'true' env: RELEASE_VERSION: ${{ env.RELEASE_VERSION }} run: | # Inject markdown release notes with sparkle:format="markdown" (Sparkle 2.9+) export NOTES=$(cat "$RELEASE_NOTES_FILE") # Insert description after the enclosure tag for this version awk ' /]*>/ && index($0, ENVIRON["RELEASE_VERSION"]) { print print " " next } { print } ' output/appcast.xml > output/appcast.xml.tmp && mv output/appcast.xml.tmp output/appcast.xml echo "Injected markdown release notes for version $RELEASE_VERSION" # ============================================================ # Upload artifacts # ============================================================ - name: Upload DMG uses: actions/upload-artifact@v4 with: name: EXO-dmg-${{ env.RELEASE_VERSION }} path: output/EXO-${{ env.RELEASE_VERSION }}.dmg - name: Upload to S3 if: env.SPARKLE_S3_BUCKET != '' env: AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} AWS_REGION: ${{ env.AWS_REGION }} SPARKLE_S3_BUCKET: ${{ env.SPARKLE_S3_BUCKET }} SPARKLE_S3_PREFIX: ${{ env.SPARKLE_S3_PREFIX }} IS_ALPHA: ${{ env.IS_ALPHA }} run: | set -euo pipefail cd output PREFIX="${SPARKLE_S3_PREFIX:-}" if [[ -n "$PREFIX" && "${PREFIX: -1}" != "/" ]]; then PREFIX="${PREFIX}/" fi DMG_NAME="EXO-${RELEASE_VERSION}.dmg" if [[ "${{ github.ref_type }}" != "tag" ]]; then aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}EXO-${GITHUB_SHA}.dmg" exit 0 fi aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}${DMG_NAME}" if [[ "$IS_ALPHA" != "true" ]]; then aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}EXO-latest.dmg" aws s3 cp appcast.xml "s3://${SPARKLE_S3_BUCKET}/${PREFIX}appcast.xml" --content-type application/xml --cache-control no-cache fi - name: Publish GitHub Release if: github.ref_type == 'tag' env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} run: | DMG_PATH="output/EXO-${RELEASE_VERSION}.dmg" if [[ "$HAS_RELEASE_NOTES" == "true" ]]; then # Update the draft release with the tag and upload DMG gh api --method PATCH "repos/${{ github.repository }}/releases/$DRAFT_RELEASE_ID" \ -f tag_name="$GITHUB_REF_NAME" \ -F draft=false gh release upload "$GITHUB_REF_NAME" "$DMG_PATH" --clobber echo "Published release $GITHUB_REF_NAME with DMG attached" else # Alpha without draft release - create one with auto-generated notes gh release create "$GITHUB_REF_NAME" "$DMG_PATH" \ --title "$GITHUB_REF_NAME" \ --generate-notes \ --prerelease echo "Created alpha release $GITHUB_REF_NAME with auto-generated notes" fi ================================================ FILE: .github/workflows/pipeline.yml ================================================ name: ci-pipeline on: push: pull_request: branches: - staging - main jobs: nix: name: Build and check (${{ matrix.system }}) runs-on: ${{ matrix.runner }} strategy: fail-fast: false matrix: include: - runner: macos-26 system: aarch64-darwin - runner: ubuntu-latest system: x86_64-linux - runner: ubuntu-24.04-arm system: aarch64-linux steps: - name: Checkout repository uses: actions/checkout@v4 with: lfs: false - uses: cachix/install-nix-action@v31 with: nix_path: nixpkgs=channel:nixos-unstable - uses: cachix/cachix-action@v14 name: Configure Cachix with: name: exo authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}" - name: Build Metal packages (macOS only) if: runner.os == 'macOS' run: | # Try to build metal-toolchain first (may succeed via cachix cache hit) if nix build .#metal-toolchain 2>/dev/null; then echo "metal-toolchain built successfully (likely cache hit)" else echo "metal-toolchain build failed, extracting from Xcode..." NAR_HASH="sha256-ayR5mXN4sZAddwKEG2OszGRF93k9ZFc7H0yi2xbylQw=" NAR_NAME="metal-toolchain-17C48.nar" # Use RUNNER_TEMP to avoid /tmp symlink issues on macOS WORK_DIR="${RUNNER_TEMP}/metal-work" mkdir -p "$WORK_DIR" # Download the Metal toolchain component xcodebuild -downloadComponent MetalToolchain # Find and mount the DMG DMG_PATH=$(find /System/Library/AssetsV2/com_apple_MobileAsset_MetalToolchain -name '*.dmg' 2>/dev/null | head -1) if [ -z "$DMG_PATH" ]; then echo "Error: Could not find Metal toolchain DMG" exit 1 fi echo "Found DMG at: $DMG_PATH" hdiutil attach "$DMG_PATH" -mountpoint "${WORK_DIR}/metal-dmg" # Copy the toolchain cp -R "${WORK_DIR}/metal-dmg/Metal.xctoolchain" "${WORK_DIR}/metal-export" hdiutil detach "${WORK_DIR}/metal-dmg" # Create NAR and add to store nix nar pack "${WORK_DIR}/metal-export" > "${WORK_DIR}/${NAR_NAME}" STORE_PATH=$(nix store add --mode flat "${WORK_DIR}/${NAR_NAME}") echo "Added NAR to store: $STORE_PATH" # Verify the hash matches ACTUAL_HASH=$(nix hash file "${WORK_DIR}/${NAR_NAME}") if [ "$ACTUAL_HASH" != "$NAR_HASH" ]; then echo "Warning: NAR hash mismatch!" echo "Expected: $NAR_HASH" echo "Actual: $ACTUAL_HASH" echo "The metal-toolchain.nix may need updating" fi # Clean up rm -rf "$WORK_DIR" # Retry the build now that NAR is in store nix build .#metal-toolchain fi # Build mlx (depends on metal-toolchain) nix build .#mlx - name: Build all Nix outputs run: | nix flake show --json | jq -r ' [ (.packages."${{ matrix.system }}" // {} | keys[] | ".#packages.${{ matrix.system }}.\(.)"), (.devShells."${{ matrix.system }}" // {} | keys[] | ".#devShells.${{ matrix.system }}.\(.)") ] | .[] ' | xargs nix build - name: Run nix flake check run: nix flake check - name: Run pytest (macOS only) if: runner.os == 'macOS' run: | # Build the test environment (requires relaxed sandbox for uv2nix on macOS) TEST_ENV=$(nix build '.#exo-test-env' --option sandbox relaxed --print-out-paths) # Run pytest outside sandbox (needs GPU access for MLX) export HOME="$RUNNER_TEMP" export EXO_TESTS=1 export EXO_DASHBOARD_DIR="$PWD/dashboard/" export EXO_RESOURCES_DIR="$PWD/resources" $TEST_ENV/bin/python -m pytest src -m "not slow" --import-mode=importlib ================================================ FILE: .gitignore ================================================ # gitingest digest.txt # python **/__pycache__ # nix .direnv/ # IDEA (PyCharm) .idea # xcode / macos *.xcuserstate *.xcuserdata *.xcuserdatad/ **/.DS_Store app/EXO/build/ dist/ # rust target/ **/*.rs.bk *.pdb # svelte dashboard/build/ dashboard/node_modules/ dashboard/.svelte-kit/ # host config snapshots hosts_*.json .swp # bench files bench/**/*.json # tmp tmp/models ================================================ FILE: .mlx_typings/.gitkeep ================================================ ================================================ FILE: .mlx_typings/mflux/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ import os if "TOKENIZERS_PARALLELISM" not in os.environ: ... ================================================ FILE: .mlx_typings/mflux/callbacks/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/callbacks/callback.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import PIL.Image import tqdm from typing import Protocol from mflux.models.common.config.config import Config class BeforeLoopCallback(Protocol): def call_before_loop( self, seed: int, prompt: str, latents: mx.array, config: Config, canny_image: PIL.Image.Image | None = ..., depth_image: PIL.Image.Image | None = ..., ) -> None: ... class InLoopCallback(Protocol): def call_in_loop( self, t: int, seed: int, prompt: str, latents: mx.array, config: Config, time_steps: tqdm, ) -> None: ... class AfterLoopCallback(Protocol): def call_after_loop( self, seed: int, prompt: str, latents: mx.array, config: Config ) -> None: ... class InterruptCallback(Protocol): def call_interrupt( self, t: int, seed: int, prompt: str, latents: mx.array, config: Config, time_steps: tqdm, ) -> None: ... ================================================ FILE: .mlx_typings/mflux/callbacks/callback_registry.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import TYPE_CHECKING from mflux.callbacks.callback import ( AfterLoopCallback, BeforeLoopCallback, InLoopCallback, InterruptCallback, ) from mflux.callbacks.generation_context import GenerationContext from mflux.models.common.config.config import Config if TYPE_CHECKING: ... class CallbackRegistry: def __init__(self) -> None: ... def register(self, callback) -> None: ... def start(self, seed: int, prompt: str, config: Config) -> GenerationContext: ... def before_loop_callbacks(self) -> list[BeforeLoopCallback]: ... def in_loop_callbacks(self) -> list[InLoopCallback]: ... def after_loop_callbacks(self) -> list[AfterLoopCallback]: ... def interrupt_callbacks(self) -> list[InterruptCallback]: ... ================================================ FILE: .mlx_typings/mflux/callbacks/generation_context.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import PIL.Image import tqdm from typing import TYPE_CHECKING from mflux.callbacks.callback_registry import CallbackRegistry from mflux.models.common.config.config import Config if TYPE_CHECKING: ... class GenerationContext: def __init__( self, registry: CallbackRegistry, seed: int, prompt: str, config: Config ) -> None: ... def before_loop( self, latents: mx.array, *, canny_image: PIL.Image.Image | None = ..., depth_image: PIL.Image.Image | None = ..., ) -> None: ... def in_loop(self, t: int, latents: mx.array, time_steps: tqdm = ...) -> None: ... def after_loop(self, latents: mx.array) -> None: ... def interruption( self, t: int, latents: mx.array, time_steps: tqdm = ... ) -> None: ... ================================================ FILE: .mlx_typings/mflux/cli/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/cli/defaults/defaults.pyi ================================================ """ This type stub file was generated by pyright. """ import os BATTERY_PERCENTAGE_STOP_LIMIT = ... CONTROLNET_STRENGTH = ... DEFAULT_DEV_FILL_GUIDANCE = ... DEFAULT_DEPTH_GUIDANCE = ... DIMENSION_STEP_PIXELS = ... GUIDANCE_SCALE = ... GUIDANCE_SCALE_KONTEXT = ... IMAGE_STRENGTH = ... MODEL_CHOICES = ... MODEL_INFERENCE_STEPS = ... QUANTIZE_CHOICES = ... if os.environ.get("MFLUX_CACHE_DIR"): MFLUX_CACHE_DIR = ... else: MFLUX_CACHE_DIR = ... MFLUX_LORA_CACHE_DIR = ... ================================================ FILE: .mlx_typings/mflux/models/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/common/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/common/cli/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/common/config/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ from mflux.models.common.config.config import Config from mflux.models.common.config.model_config import ModelConfig __all__ = ["Config", "ModelConfig"] ================================================ FILE: .mlx_typings/mflux/models/common/config/config.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from pathlib import Path from typing import Any from tqdm import tqdm from mflux.models.common.config.model_config import ModelConfig logger = ... class Config: def __init__( self, model_config: ModelConfig, num_inference_steps: int = ..., height: int = ..., width: int = ..., guidance: float = ..., image_path: Path | str | None = ..., image_strength: float | None = ..., depth_image_path: Path | str | None = ..., redux_image_paths: list[Path | str] | None = ..., redux_image_strengths: list[float] | None = ..., masked_image_path: Path | str | None = ..., controlnet_strength: float | None = ..., scheduler: str = ..., ) -> None: ... @property def height(self) -> int: ... @property def width(self) -> int: ... @width.setter def width(self, value): # -> None: ... @property def image_seq_len(self) -> int: ... @property def guidance(self) -> float: ... @property def num_inference_steps(self) -> int: ... @property def precision(self) -> mx.Dtype: ... @property def num_train_steps(self) -> int: ... @property def image_path(self) -> Path | None: ... @property def image_strength(self) -> float | None: ... @property def depth_image_path(self) -> Path | None: ... @property def redux_image_paths(self) -> list[Path] | None: ... @property def redux_image_strengths(self) -> list[float] | None: ... @property def masked_image_path(self) -> Path | None: ... @property def init_time_step(self) -> int: ... @property def time_steps(self) -> tqdm: ... @property def controlnet_strength(self) -> float | None: ... @property def scheduler(self) -> Any: ... ================================================ FILE: .mlx_typings/mflux/models/common/config/model_config.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from functools import lru_cache from typing import Literal class ModelConfig: precision: mx.Dtype = ... def __init__( self, priority: int, aliases: list[str], model_name: str, base_model: str | None, controlnet_model: str | None, custom_transformer_model: str | None, num_train_steps: int | None, max_sequence_length: int | None, supports_guidance: bool | None, requires_sigma_shift: bool | None, transformer_overrides: dict | None = ..., ) -> None: ... @staticmethod @lru_cache def dev() -> ModelConfig: ... @staticmethod @lru_cache def schnell() -> ModelConfig: ... @staticmethod @lru_cache def dev_kontext() -> ModelConfig: ... @staticmethod @lru_cache def dev_fill() -> ModelConfig: ... @staticmethod @lru_cache def dev_redux() -> ModelConfig: ... @staticmethod @lru_cache def dev_depth() -> ModelConfig: ... @staticmethod @lru_cache def dev_controlnet_canny() -> ModelConfig: ... @staticmethod @lru_cache def schnell_controlnet_canny() -> ModelConfig: ... @staticmethod @lru_cache def dev_controlnet_upscaler() -> ModelConfig: ... @staticmethod @lru_cache def dev_fill_catvton() -> ModelConfig: ... @staticmethod @lru_cache def krea_dev() -> ModelConfig: ... @staticmethod @lru_cache def flux2_klein_4b() -> ModelConfig: ... @staticmethod @lru_cache def flux2_klein_9b() -> ModelConfig: ... @staticmethod @lru_cache def qwen_image() -> ModelConfig: ... @staticmethod @lru_cache def qwen_image_edit() -> ModelConfig: ... @staticmethod @lru_cache def fibo() -> ModelConfig: ... @staticmethod @lru_cache def z_image_turbo() -> ModelConfig: ... @staticmethod @lru_cache def seedvr2_3b() -> ModelConfig: ... def x_embedder_input_dim(self) -> int: ... def is_canny(self) -> bool: ... @staticmethod def from_name( model_name: str, base_model: Literal["dev", "schnell", "krea-dev"] | None = ... ) -> ModelConfig: ... AVAILABLE_MODELS = ... ================================================ FILE: .mlx_typings/mflux/models/common/latent_creator/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/common/latent_creator/latent_creator.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from pathlib import Path from typing import TYPE_CHECKING, TypeAlias from mlx import nn from mflux.models.common.vae.tiling_config import TilingConfig from mflux.models.fibo.latent_creator.fibo_latent_creator import FiboLatentCreator from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator from mflux.models.z_image.latent_creator.z_image_latent_creator import ( ZImageLatentCreator, ) if TYPE_CHECKING: LatentCreatorType: TypeAlias = type[ FiboLatentCreator | FluxLatentCreator | QwenLatentCreator | ZImageLatentCreator ] class Img2Img: def __init__( self, vae: nn.Module, latent_creator: LatentCreatorType, sigmas: mx.array, init_time_step: int, image_path: str | Path | None, tiling_config: TilingConfig | None = ..., ) -> None: ... class LatentCreator: @staticmethod def create_for_txt2img_or_img2img( seed: int, height: int, width: int, img2img: Img2Img ) -> mx.array: ... @staticmethod def encode_image( vae: nn.Module, image_path: str | Path, height: int, width: int, tiling_config: TilingConfig | None = ..., ) -> mx.array: ... @staticmethod def add_noise_by_interpolation( clean: mx.array, noise: mx.array, sigma: float ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/common/lora/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/common/lora/layer/fused_linear_lora_layer.pyi ================================================ """ This type stub file was generated by pyright. """ from mlx import nn from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear class FusedLoRALinear(nn.Module): def __init__( self, base_linear: nn.Linear | nn.QuantizedLinear, loras: list[LoRALinear] ) -> None: ... def __call__(self, x): # -> array: ... ================================================ FILE: .mlx_typings/mflux/models/common/lora/layer/linear_lora_layer.pyi ================================================ """ This type stub file was generated by pyright. """ from mlx import nn class LoRALinear(nn.Module): @staticmethod def from_linear( linear: nn.Linear | nn.QuantizedLinear, r: int = ..., scale: float = ... ): # -> LoRALinear: ... def __init__( self, input_dims: int, output_dims: int, r: int = ..., scale: float = ..., bias: bool = ..., ) -> None: ... def __call__(self, x): # -> array: ... ================================================ FILE: .mlx_typings/mflux/models/common/lora/mapping/lora_loader.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn from collections.abc import Callable from dataclasses import dataclass from mflux.models.common.lora.mapping.lora_mapping import LoRATarget @dataclass class PatternMatch: source_pattern: str target_path: str matrix_name: str transpose: bool transform: Callable[[mx.array], mx.array] | None = ... class LoRALoader: @staticmethod def load_and_apply_lora( lora_mapping: list[LoRATarget], transformer: nn.Module, lora_paths: list[str] | None = ..., lora_scales: list[float] | None = ..., ) -> tuple[list[str], list[float]]: ... ================================================ FILE: .mlx_typings/mflux/models/common/lora/mapping/lora_mapping.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from collections.abc import Callable from dataclasses import dataclass from typing import List, Protocol @dataclass class LoRATarget: model_path: str possible_up_patterns: List[str] possible_down_patterns: List[str] possible_alpha_patterns: List[str] = ... up_transform: Callable[[mx.array], mx.array] | None = ... down_transform: Callable[[mx.array], mx.array] | None = ... class LoRAMapping(Protocol): @staticmethod def get_mapping() -> List[LoRATarget]: ... ================================================ FILE: .mlx_typings/mflux/models/common/lora/mapping/lora_saver.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.nn as nn class LoRASaver: @staticmethod def bake_and_strip_lora(module: nn.Module) -> nn.Module: ... ================================================ FILE: .mlx_typings/mflux/models/common/lora/mapping/lora_transforms.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx class LoraTransforms: @staticmethod def split_q_up(tensor: mx.array) -> mx.array: ... @staticmethod def split_k_up(tensor: mx.array) -> mx.array: ... @staticmethod def split_v_up(tensor: mx.array) -> mx.array: ... @staticmethod def split_q_down(tensor: mx.array) -> mx.array: ... @staticmethod def split_k_down(tensor: mx.array) -> mx.array: ... @staticmethod def split_v_down(tensor: mx.array) -> mx.array: ... @staticmethod def split_single_q_up(tensor: mx.array) -> mx.array: ... @staticmethod def split_single_k_up(tensor: mx.array) -> mx.array: ... @staticmethod def split_single_v_up(tensor: mx.array) -> mx.array: ... @staticmethod def split_single_mlp_up(tensor: mx.array) -> mx.array: ... @staticmethod def split_single_q_down(tensor: mx.array) -> mx.array: ... @staticmethod def split_single_k_down(tensor: mx.array) -> mx.array: ... @staticmethod def split_single_v_down(tensor: mx.array) -> mx.array: ... @staticmethod def split_single_mlp_down(tensor: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/common/resolution/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ from mflux.models.common.resolution.config_resolution import ConfigResolution from mflux.models.common.resolution.lora_resolution import LoraResolution from mflux.models.common.resolution.path_resolution import PathResolution from mflux.models.common.resolution.quantization_resolution import ( QuantizationResolution, ) __all__ = [ "ConfigResolution", "LoraResolution", "PathResolution", "QuantizationResolution", ] ================================================ FILE: .mlx_typings/mflux/models/common/resolution/actions.pyi ================================================ """ This type stub file was generated by pyright. """ from enum import Enum from typing import NamedTuple class QuantizationAction(Enum): NONE = ... STORED = ... REQUESTED = ... class PathAction(Enum): LOCAL = ... HUGGINGFACE_CACHED = ... HUGGINGFACE = ... ERROR = ... class LoraAction(Enum): LOCAL = ... REGISTRY = ... HUGGINGFACE_COLLECTION_CACHED = ... HUGGINGFACE_COLLECTION = ... HUGGINGFACE_REPO_CACHED = ... HUGGINGFACE_REPO = ... ERROR = ... class ConfigAction(Enum): EXACT_MATCH = ... EXPLICIT_BASE = ... INFER_SUBSTRING = ... ERROR = ... class Rule(NamedTuple): priority: int name: str check: str action: QuantizationAction | PathAction | LoraAction | ConfigAction ... ================================================ FILE: .mlx_typings/mflux/models/common/resolution/config_resolution.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import TYPE_CHECKING from mflux.models.common.config.model_config import ModelConfig if TYPE_CHECKING: ... logger = ... class ConfigResolution: RULES = ... @staticmethod def resolve(model_name: str, base_model: str | None = ...) -> ModelConfig: ... ================================================ FILE: .mlx_typings/mflux/models/common/resolution/lora_resolution.pyi ================================================ """ This type stub file was generated by pyright. """ from pathlib import Path logger = ... class LoraResolution: RULES = ... _registry: dict[str, Path] = ... @staticmethod def resolve(path: str) -> str: ... @staticmethod def resolve_paths(paths: list[str] | None) -> list[str]: ... @staticmethod def resolve_scales(scales: list[float] | None, num_paths: int) -> list[float]: ... @staticmethod def get_registry() -> dict[str, Path]: ... @staticmethod def discover_files(library_paths: list[Path]) -> dict[str, Path]: ... ================================================ FILE: .mlx_typings/mflux/models/common/resolution/path_resolution.pyi ================================================ """ This type stub file was generated by pyright. """ from pathlib import Path logger = ... class PathResolution: RULES = ... @staticmethod def resolve(path: str | None, patterns: list[str] | None = ...) -> Path | None: ... ================================================ FILE: .mlx_typings/mflux/models/common/resolution/quantization_resolution.pyi ================================================ """ This type stub file was generated by pyright. """ logger = ... class QuantizationResolution: RULES = ... @staticmethod def resolve( stored: int | None, requested: int | None ) -> tuple[int | None, str | None]: ... ================================================ FILE: .mlx_typings/mflux/models/common/schedulers/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ from .flow_match_euler_discrete_scheduler import FlowMatchEulerDiscreteScheduler from .linear_scheduler import LinearScheduler from .seedvr2_euler_scheduler import SeedVR2EulerScheduler __all__ = [ "LinearScheduler", "FlowMatchEulerDiscreteScheduler", "SeedVR2EulerScheduler", ] class SchedulerModuleNotFound(ValueError): ... class SchedulerClassNotFound(ValueError): ... class InvalidSchedulerType(TypeError): ... SCHEDULER_REGISTRY = ... def register_contrib(scheduler_object, scheduler_name=...): # -> None: ... def try_import_external_scheduler( scheduler_object_path: str, ): # -> type[BaseScheduler]: ... ================================================ FILE: .mlx_typings/mflux/models/common/schedulers/base_scheduler.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from abc import ABC, abstractmethod class BaseScheduler(ABC): @property @abstractmethod def sigmas(self) -> mx.array: ... @abstractmethod def step( self, noise: mx.array, timestep: int, latents: mx.array, **kwargs ) -> mx.array: ... def scale_model_input(self, latents: mx.array, t: int) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/common/schedulers/flow_match_euler_discrete_scheduler.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from typing import TYPE_CHECKING from mflux.models.common.config.config import Config from mflux.models.common.schedulers.base_scheduler import BaseScheduler if TYPE_CHECKING: ... class FlowMatchEulerDiscreteScheduler(BaseScheduler): def __init__(self, config: Config) -> None: ... @property def sigmas(self) -> mx.array: ... @property def timesteps(self) -> mx.array: ... def set_image_seq_len(self, image_seq_len: int) -> None: ... @staticmethod def get_timesteps_and_sigmas( image_seq_len: int, num_inference_steps: int, num_train_timesteps: int = ... ) -> tuple[mx.array, mx.array]: ... def step( self, noise: mx.array, timestep: int, latents: mx.array, **kwargs ) -> mx.array: ... def scale_model_input(self, latents: mx.array, t: int) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/common/schedulers/linear_scheduler.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from typing import TYPE_CHECKING from mflux.models.common.config.config import Config from mflux.models.common.schedulers.base_scheduler import BaseScheduler if TYPE_CHECKING: ... class LinearScheduler(BaseScheduler): def __init__(self, config: Config) -> None: ... @property def sigmas(self) -> mx.array: ... @property def timesteps(self) -> mx.array: ... def step( self, noise: mx.array, timestep: int, latents: mx.array, **kwargs ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/common/schedulers/seedvr2_euler_scheduler.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from typing import TYPE_CHECKING from mflux.models.common.config.config import Config from mflux.models.common.schedulers.base_scheduler import BaseScheduler if TYPE_CHECKING: ... class SeedVR2EulerScheduler(BaseScheduler): def __init__(self, config: Config) -> None: ... @property def timesteps(self) -> mx.array: ... @property def sigmas(self) -> mx.array: ... def step( self, noise: mx.array, timestep: int, latents: mx.array, **kwargs ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/common/tokenizer/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ from mflux.models.common.tokenizer.tokenizer import ( BaseTokenizer, LanguageTokenizer, Tokenizer, VisionLanguageTokenizer, ) from mflux.models.common.tokenizer.tokenizer_loader import TokenizerLoader from mflux.models.common.tokenizer.tokenizer_output import TokenizerOutput """ This type stub file was generated by pyright. """ __all__ = [ "Tokenizer", "BaseTokenizer", "LanguageTokenizer", "VisionLanguageTokenizer", "TokenizerLoader", "TokenizerOutput", ] ================================================ FILE: .mlx_typings/mflux/models/common/tokenizer/tokenizer.pyi ================================================ """ This type stub file was generated by pyright. """ from abc import ABC, abstractmethod from typing import Protocol, runtime_checkable from PIL import Image from transformers import PreTrainedTokenizer from mflux.models.common.tokenizer.tokenizer_output import TokenizerOutput """ This type stub file was generated by pyright. """ @runtime_checkable class Tokenizer(Protocol): tokenizer: PreTrainedTokenizer def tokenize( self, prompt: str | list[str], images: list[Image.Image] | None = ..., max_length: int | None = ..., **kwargs, ) -> TokenizerOutput: ... class BaseTokenizer(ABC): def __init__( self, tokenizer: PreTrainedTokenizer, max_length: int = ... ) -> None: ... @abstractmethod def tokenize( self, prompt: str | list[str], images: list[Image.Image] | None = ..., max_length: int | None = ..., **kwargs, ) -> TokenizerOutput: ... class LanguageTokenizer(BaseTokenizer): def __init__( self, tokenizer: PreTrainedTokenizer, max_length: int = ..., padding: str = ..., return_attention_mask: bool = ..., template: str | None = ..., use_chat_template: bool = ..., chat_template_kwargs: dict | None = ..., add_special_tokens: bool = ..., ) -> None: ... def tokenize( self, prompt: str | list[str], images: list[Image.Image] | None = ..., max_length: int | None = ..., **kwargs, ) -> TokenizerOutput: ... class VisionLanguageTokenizer(BaseTokenizer): def __init__( self, tokenizer: PreTrainedTokenizer, processor, max_length: int = ..., template: str | None = ..., image_token: str = ..., ) -> None: ... def tokenize( self, prompt: str | list[str], images: list[Image.Image] | None = ..., max_length: int | None = ..., **kwargs, ) -> TokenizerOutput: ... ================================================ FILE: .mlx_typings/mflux/models/common/tokenizer/tokenizer_loader.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import TYPE_CHECKING from mflux.models.common.tokenizer.tokenizer import BaseTokenizer from mflux.models.common.weights.loading.weight_definition import TokenizerDefinition """ This type stub file was generated by pyright. """ if TYPE_CHECKING: ... class TokenizerLoader: @staticmethod def load(definition: TokenizerDefinition, model_path: str) -> BaseTokenizer: ... @staticmethod def load_all( definitions: list[TokenizerDefinition], model_path: str, max_length_overrides: dict[str, int] | None = ..., ) -> dict[str, BaseTokenizer]: ... ================================================ FILE: .mlx_typings/mflux/models/common/tokenizer/tokenizer_output.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from dataclasses import dataclass """ This type stub file was generated by pyright. """ @dataclass class TokenizerOutput: input_ids: mx.array attention_mask: mx.array pixel_values: mx.array | None = ... image_grid_thw: mx.array | None = ... ================================================ FILE: .mlx_typings/mflux/models/common/vae/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ from mflux.models.common.vae.tiling_config import TilingConfig from mflux.models.common.vae.vae_tiler import VAETiler __all__ = ["TilingConfig", "VAETiler"] ================================================ FILE: .mlx_typings/mflux/models/common/vae/tiling_config.pyi ================================================ """ This type stub file was generated by pyright. """ from dataclasses import dataclass @dataclass(frozen=True, slots=True) class TilingConfig: vae_decode_tiles_per_dim: int | None = ... vae_decode_overlap: int = ... vae_encode_tiled: bool = ... vae_encode_tile_size: int = ... vae_encode_tile_overlap: int = ... ================================================ FILE: .mlx_typings/mflux/models/common/vae/vae_tiler.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from typing import Callable class VAETiler: @staticmethod def encode_image_tiled( *, image: mx.array, encode_fn: Callable[[mx.array], mx.array], latent_channels: int, tile_size: tuple[int, int] = ..., tile_overlap: tuple[int, int] = ..., spatial_scale: int = ..., ) -> mx.array: ... @staticmethod def decode_image_tiled( *, latent: mx.array, decode_fn: Callable[[mx.array], mx.array], tile_size: tuple[int, int] = ..., tile_overlap: tuple[int, int] = ..., spatial_scale: int = ..., ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/common/vae/vae_util.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn from mflux.models.common.vae.tiling_config import TilingConfig class VAEUtil: @staticmethod def encode( vae: nn.Module, image: mx.array, tiling_config: TilingConfig | None = ... ) -> mx.array: ... @staticmethod def decode( vae: nn.Module, latent: mx.array, tiling_config: TilingConfig | None = ... ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/common/weights/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ from mflux.models.common.weights.loading.loaded_weights import LoadedWeights, MetaData from mflux.models.common.weights.loading.weight_applier import WeightApplier from mflux.models.common.weights.loading.weight_definition import ComponentDefinition from mflux.models.common.weights.loading.weight_loader import WeightLoader from mflux.models.common.weights.saving.model_saver import ModelSaver __all__ = [ "ComponentDefinition", "LoadedWeights", "MetaData", "ModelSaver", "WeightApplier", "WeightLoader", ] ================================================ FILE: .mlx_typings/mflux/models/common/weights/loading/loaded_weights.pyi ================================================ """ This type stub file was generated by pyright. """ from dataclasses import dataclass @dataclass class MetaData: quantization_level: int | None = ... mflux_version: str | None = ... @dataclass class LoadedWeights: components: dict[str, dict] meta_data: MetaData def __getattr__(self, name: str) -> dict | None: ... def num_transformer_blocks(self, component_name: str = ...) -> int: ... def num_single_transformer_blocks(self, component_name: str = ...) -> int: ... ================================================ FILE: .mlx_typings/mflux/models/common/weights/loading/weight_applier.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.nn as nn from typing import TYPE_CHECKING from mflux.models.common.weights.loading.loaded_weights import LoadedWeights from mflux.models.common.weights.loading.weight_definition import ( ComponentDefinition, WeightDefinitionType, ) if TYPE_CHECKING: ... class WeightApplier: @staticmethod def apply_and_quantize_single( weights: LoadedWeights, model: nn.Module, component: ComponentDefinition, quantize_arg: int | None, quantization_predicate=..., ) -> int | None: ... @staticmethod def apply_and_quantize( weights: LoadedWeights, models: dict[str, nn.Module], quantize_arg: int | None, weight_definition: WeightDefinitionType, ) -> int | None: ... ================================================ FILE: .mlx_typings/mflux/models/common/weights/loading/weight_definition.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from dataclasses import dataclass from typing import Callable, List, TYPE_CHECKING, TypeAlias from mflux.models.common.weights.mapping.weight_mapping import WeightTarget from mflux.models.common.tokenizer.tokenizer import BaseTokenizer from mflux.models.depth_pro.weights.depth_pro_weight_definition import ( DepthProWeightDefinition, ) from mflux.models.fibo.weights.fibo_weight_definition import FIBOWeightDefinition from mflux.models.fibo_vlm.weights.fibo_vlm_weight_definition import ( FIBOVLMWeightDefinition, ) from mflux.models.flux.weights.flux_weight_definition import FluxWeightDefinition from mflux.models.qwen.weights.qwen_weight_definition import QwenWeightDefinition from mflux.models.seedvr2.weights.seedvr2_weight_definition import ( SeedVR2WeightDefinition, ) from mflux.models.z_image.weights.z_image_weight_definition import ( ZImageWeightDefinition, ) """ This type stub file was generated by pyright. """ if TYPE_CHECKING: WeightDefinitionType: TypeAlias = type[ FluxWeightDefinition | FIBOWeightDefinition | FIBOVLMWeightDefinition | QwenWeightDefinition | ZImageWeightDefinition | SeedVR2WeightDefinition | DepthProWeightDefinition ] @dataclass class ComponentDefinition: name: str hf_subdir: str mapping_getter: Callable[[], List[WeightTarget]] | None = ... model_attr: str | None = ... num_blocks: int | None = ... num_layers: int | None = ... loading_mode: str = ... precision: mx.Dtype | None = ... skip_quantization: bool = ... bulk_transform: Callable[[mx.array], mx.array] | None = ... weight_subkey: str | None = ... download_url: str | None = ... weight_prefix_filters: List[str] | None = ... weight_files: List[str] | None = ... @dataclass class TokenizerDefinition: name: str hf_subdir: str tokenizer_class: str = ... fallback_subdirs: List[str] | None = ... download_patterns: List[str] | None = ... encoder_class: type[BaseTokenizer] | None = ... max_length: int = ... padding: str = ... template: str | None = ... use_chat_template: bool = ... chat_template_kwargs: dict | None = ... add_special_tokens: bool = ... processor_class: type | None = ... image_token: str = ... chat_template: str | None = ... ================================================ FILE: .mlx_typings/mflux/models/common/weights/loading/weight_loader.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import TYPE_CHECKING from mflux.models.common.weights.loading.loaded_weights import LoadedWeights from mflux.models.common.weights.loading.weight_definition import ( ComponentDefinition, WeightDefinitionType, ) if TYPE_CHECKING: ... logger = ... class WeightLoader: @staticmethod def load_single( component: ComponentDefinition, repo_id: str, file_pattern: str = ... ) -> LoadedWeights: ... @staticmethod def load( weight_definition: WeightDefinitionType, model_path: str | None = ... ) -> LoadedWeights: ... ================================================ FILE: .mlx_typings/mflux/models/common/weights/mapping/weight_mapper.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from typing import Dict, List, Optional from mflux.models.common.weights.mapping.weight_mapping import WeightTarget class WeightMapper: @staticmethod def apply_mapping( hf_weights: Dict[str, mx.array], mapping: List[WeightTarget], num_blocks: Optional[int] = ..., num_layers: Optional[int] = ..., ) -> Dict: ... ================================================ FILE: .mlx_typings/mflux/models/common/weights/mapping/weight_mapping.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from dataclasses import dataclass from typing import Callable, List, Optional, Protocol """ This type stub file was generated by pyright. """ @dataclass class WeightTarget: to_pattern: str from_pattern: List[str] transform: Optional[Callable[[mx.array], mx.array]] = ... required: bool = ... max_blocks: Optional[int] = ... class WeightMapping(Protocol): @staticmethod def get_mapping() -> List[WeightTarget]: ... ================================================ FILE: .mlx_typings/mflux/models/common/weights/mapping/weight_transforms.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx class WeightTransforms: @staticmethod def reshape_gamma_to_1d(tensor: mx.array) -> mx.array: ... @staticmethod def transpose_patch_embed(tensor: mx.array) -> mx.array: ... @staticmethod def transpose_conv3d_weight(tensor: mx.array) -> mx.array: ... @staticmethod def transpose_conv2d_weight(tensor: mx.array) -> mx.array: ... @staticmethod def transpose_conv_transpose2d_weight(tensor: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/common/weights/saving/model_saver.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Any, TYPE_CHECKING from mflux.models.common.weights.loading.weight_definition import WeightDefinitionType if TYPE_CHECKING: ... class ModelSaver: @staticmethod def save_model( model: Any, bits: int, base_path: str, weight_definition: WeightDefinitionType ) -> None: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/depth_pro_initializer.pyi ================================================ """ This type stub file was generated by pyright. """ from mflux.models.depth_pro.model.depth_pro_model import DepthProModel class DepthProInitializer: @staticmethod def init(model: DepthProModel, quantize: int | None = ...) -> None: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/decoder/feature_fusion_block_2d.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class FeatureFusionBlock2d(nn.Module): def __init__(self, num_features: int, deconv: bool = ...) -> None: ... def __call__(self, x0: mx.array, x1: mx.array | None = ...) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/decoder/multires_conv_decoder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class MultiresConvDecoder(nn.Module): def __init__(self) -> None: ... def __call__( self, x0_latent: mx.array, x1_latent: mx.array, x0_features: mx.array, x1_features: mx.array, x_global_features: mx.array, ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/decoder/residual_block.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class ResidualBlock(nn.Module): def __init__(self, num_features: int) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/depth_pro.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from dataclasses import dataclass from pathlib import Path from PIL import Image @dataclass class DepthResult: depth_image: Image.Image depth_array: mx.array min_depth: float max_depth: float ... class DepthPro: def __init__(self, quantize: int | None = ...) -> None: ... def create_depth_map(self, image_path: str | Path) -> DepthResult: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/depth_pro_model.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class DepthProModel(nn.Module): def __init__(self) -> None: ... def __call__( self, x0: mx.array, x1: mx.array, x2: mx.array ) -> tuple[mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/depth_pro_util.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class DepthProUtil: @staticmethod def split(x: mx.array, overlap_ratio: float = ...) -> mx.array: ... @staticmethod def interpolate(x: mx.array, size=..., scale_factor=...): # -> array: ... @staticmethod def apply_conv(x: mx.array, conv_module: nn.Module) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/dino_v2/attention.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class Attention(nn.Module): def __init__( self, dim: int = ..., head_dim: int = ..., num_heads: int = ... ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/dino_v2/dino_vision_transformer.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class DinoVisionTransformer(nn.Module): def __init__(self) -> None: ... def __call__(self, x: mx.array) -> tuple[mx.array, mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/dino_v2/layer_scale.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class LayerScale(nn.Module): def __init__(self, dims: int, init_values: float = ...) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/dino_v2/mlp.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class MLP(nn.Module): def __init__(self) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/dino_v2/patch_embed.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class PatchEmbed(nn.Module): def __init__(self) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/dino_v2/transformer_block.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class TransformerBlock(nn.Module): def __init__(self) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/encoder/depth_pro_encoder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class DepthProEncoder(nn.Module): def __init__(self) -> None: ... def __call__( self, x0: mx.array, x1: mx.array, x2: mx.array ) -> tuple[mx.array, mx.array, mx.array, mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/encoder/upsample_block.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class UpSampleBlock(nn.Module): def __init__( self, dim_in: int = ..., dim_int: int = ..., dim_out: int = ..., upsample_layers: int = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/model/head/fov_head.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class FOVHead(nn.Module): def __init__(self) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/weights/depth_pro_weight_definition.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import List from mflux.models.common.weights.loading.weight_definition import ( ComponentDefinition, TokenizerDefinition, ) """ This type stub file was generated by pyright. """ class DepthProWeightDefinition: @staticmethod def get_components() -> List[ComponentDefinition]: ... @staticmethod def get_tokenizers() -> List[TokenizerDefinition]: ... @staticmethod def get_download_patterns() -> List[str]: ... @staticmethod def quantization_predicate(path: str, module) -> bool: ... ================================================ FILE: .mlx_typings/mflux/models/depth_pro/weights/depth_pro_weight_mapping.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import List from mflux.models.common.weights.mapping.weight_mapping import ( WeightMapping, WeightTarget, ) class DepthProWeightMapping(WeightMapping): @staticmethod def get_mapping() -> List[WeightTarget]: ... ================================================ FILE: .mlx_typings/mflux/models/fibo/latent_creator/fibo_latent_creator.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx class FiboLatentCreator: @staticmethod def create_noise(seed: int, height: int, width: int) -> mx.array: ... @staticmethod def pack_latents(latents: mx.array, height: int, width: int) -> mx.array: ... @staticmethod def unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/fibo/weights/fibo_weight_definition.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import List from mflux.models.common.weights.loading.weight_definition import ( ComponentDefinition, TokenizerDefinition, ) """ This type stub file was generated by pyright. """ class FIBOWeightDefinition: @staticmethod def get_components() -> List[ComponentDefinition]: ... @staticmethod def get_tokenizers() -> List[TokenizerDefinition]: ... @staticmethod def get_download_patterns() -> List[str]: ... @staticmethod def quantization_predicate(path: str, module) -> bool: ... ================================================ FILE: .mlx_typings/mflux/models/fibo/weights/fibo_weight_mapping.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import List from mflux.models.common.weights.mapping.weight_mapping import ( WeightMapping, WeightTarget, ) class FIBOWeightMapping(WeightMapping): @staticmethod def get_transformer_mapping() -> List[WeightTarget]: ... @staticmethod def get_text_encoder_mapping() -> List[WeightTarget]: ... @staticmethod def get_vae_mapping() -> List[WeightTarget]: ... ================================================ FILE: .mlx_typings/mflux/models/fibo_vlm/tokenizer/qwen2vl_image_processor.pyi ================================================ """ This type stub file was generated by pyright. """ from mflux.models.qwen.tokenizer.qwen_image_processor import QwenImageProcessor class Qwen2VLImageProcessor(QwenImageProcessor): def __init__(self) -> None: ... ================================================ FILE: .mlx_typings/mflux/models/fibo_vlm/tokenizer/qwen2vl_processor.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Optional, Union from PIL import Image class Qwen2VLProcessor: def __init__(self, tokenizer) -> None: ... def apply_chat_template( self, messages, tokenize: bool = ..., add_generation_prompt: bool = ..., return_tensors: Optional[str] = ..., return_dict: bool = ..., **kwargs, ): # -> dict[Any, Any]: ... def __call__( self, text: Optional[Union[str, list[str]]] = ..., images: Optional[Union[Image.Image, list[Image.Image]]] = ..., padding: bool = ..., return_tensors: Optional[str] = ..., **kwargs, ): # -> dict[Any, Any]: ... ================================================ FILE: .mlx_typings/mflux/models/fibo_vlm/weights/fibo_vlm_weight_definition.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import List from mflux.models.common.weights.loading.weight_definition import ( ComponentDefinition, TokenizerDefinition, ) """ This type stub file was generated by pyright. """ QWEN2VL_CHAT_TEMPLATE = ... class FIBOVLMWeightDefinition: @staticmethod def get_components() -> List[ComponentDefinition]: ... @staticmethod def get_tokenizers() -> List[TokenizerDefinition]: ... @staticmethod def get_download_patterns() -> List[str]: ... @staticmethod def quantization_predicate(path: str, module) -> bool: ... ================================================ FILE: .mlx_typings/mflux/models/fibo_vlm/weights/fibo_vlm_weight_mapping.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import List from mflux.models.common.weights.mapping.weight_mapping import ( WeightMapping, WeightTarget, ) class FIBOVLMWeightMapping(WeightMapping): @staticmethod def get_vlm_decoder_mapping(num_layers: int = ...) -> List[WeightTarget]: ... @staticmethod def get_vlm_visual_mapping(depth: int = ...) -> List[WeightTarget]: ... ================================================ FILE: .mlx_typings/mflux/models/flux/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/flux/cli/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/flux/flux_initializer.pyi ================================================ """ This type stub file was generated by pyright. """ from mflux.models.common.config import ModelConfig class FluxInitializer: @staticmethod def init( model, model_config: ModelConfig, quantize: int | None, model_path: str | None = ..., lora_paths: list[str] | None = ..., lora_scales: list[float] | None = ..., custom_transformer=..., ) -> None: ... @staticmethod def init_depth( model, model_config: ModelConfig, quantize: int | None, model_path: str | None = ..., lora_paths: list[str] | None = ..., lora_scales: list[float] | None = ..., ) -> None: ... @staticmethod def init_redux( model, model_config: ModelConfig, quantize: int | None, model_path: str | None = ..., lora_paths: list[str] | None = ..., lora_scales: list[float] | None = ..., ) -> None: ... @staticmethod def init_controlnet( model, model_config: ModelConfig, quantize: int | None, model_path: str | None = ..., lora_paths: list[str] | None = ..., lora_scales: list[float] | None = ..., ) -> None: ... @staticmethod def init_concept( model, model_config: ModelConfig, quantize: int | None, model_path: str | None = ..., lora_paths: list[str] | None = ..., lora_scales: list[float] | None = ..., ) -> None: ... ================================================ FILE: .mlx_typings/mflux/models/flux/latent_creator/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/flux/latent_creator/flux_latent_creator.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx """ This type stub file was generated by pyright. """ class FluxLatentCreator: @staticmethod def create_noise(seed: int, height: int, width: int) -> mx.array: ... @staticmethod def pack_latents( latents: mx.array, height: int, width: int, num_channels_latents: int = ... ) -> mx.array: ... @staticmethod def unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_embeddings.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class CLIPEmbeddings(nn.Module): def __init__(self, dims: int) -> None: ... def __call__(self, tokens: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_encoder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn """ This type stub file was generated by pyright. """ class CLIPEncoder(nn.Module): def __init__(self) -> None: ... def __call__(self, tokens: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_encoder_layer.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class CLIPEncoderLayer(nn.Module): def __init__(self, layer: int) -> None: ... def __call__( self, hidden_states: mx.array, causal_attention_mask: mx.array ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_mlp.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class CLIPMLP(nn.Module): def __init__(self) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... @staticmethod def quick_gelu(input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_sdpa_attention.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class CLIPSdpaAttention(nn.Module): head_dimension = ... batch_size = ... num_heads = ... def __init__(self) -> None: ... def __call__( self, hidden_states: mx.array, causal_attention_mask: mx.array ) -> mx.array: ... @staticmethod def reshape_and_transpose(x, batch_size, num_heads, head_dim): # -> array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_text_model.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class CLIPTextModel(nn.Module): def __init__(self, dims: int, num_encoder_layers: int) -> None: ... def __call__(self, tokens: mx.array) -> tuple[mx.array, mx.array]: ... @staticmethod def create_causal_attention_mask(input_shape: tuple) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/clip_encoder/encoder_clip.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class EncoderCLIP(nn.Module): def __init__(self, num_encoder_layers: int) -> None: ... def __call__( self, tokens: mx.array, causal_attention_mask: mx.array ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/prompt_encoder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mflux.models.common.tokenizer import Tokenizer from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import ( CLIPEncoder, ) from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5Encoder """ This type stub file was generated by pyright. """ class PromptEncoder: @staticmethod def encode_prompt( prompt: str, prompt_cache: dict[str, tuple[mx.array, mx.array]], t5_tokenizer: Tokenizer, clip_tokenizer: Tokenizer, t5_text_encoder: T5Encoder, clip_text_encoder: CLIPEncoder, ) -> tuple[mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_attention.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class T5Attention(nn.Module): def __init__(self) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_block.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class T5Block(nn.Module): def __init__(self, layer: int) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_dense_relu_dense.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class T5DenseReluDense(nn.Module): def __init__(self) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... @staticmethod def new_gelu(input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_encoder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn """ This type stub file was generated by pyright. """ class T5Encoder(nn.Module): def __init__(self) -> None: ... def __call__(self, tokens: mx.array): ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_feed_forward.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class T5FeedForward(nn.Module): def __init__(self) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_layer_norm.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class T5LayerNorm(nn.Module): def __init__(self) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_text_encoder/t5_encoder/t5_self_attention.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class T5SelfAttention(nn.Module): def __init__(self) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... @staticmethod def shape(states): # -> array: ... @staticmethod def un_shape(states): # -> array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/ada_layer_norm_continuous.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class AdaLayerNormContinuous(nn.Module): def __init__(self, embedding_dim: int, conditioning_embedding_dim: int) -> None: ... def __call__(self, x: mx.array, text_embeddings: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/ada_layer_norm_zero.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class AdaLayerNormZero(nn.Module): def __init__(self) -> None: ... def __call__( self, hidden_states: mx.array, text_embeddings: mx.array ) -> tuple[mx.array, mx.array, mx.array, mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/ada_layer_norm_zero_single.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class AdaLayerNormZeroSingle(nn.Module): def __init__(self) -> None: ... def __call__( self, hidden_states: mx.array, text_embeddings: mx.array ) -> tuple[mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/common/attention_utils.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class AttentionUtils: @staticmethod def process_qkv( hidden_states: mx.array, to_q: nn.Linear, to_k: nn.Linear, to_v: nn.Linear, norm_q: nn.RMSNorm, norm_k: nn.RMSNorm, num_heads: int, head_dim: int, ) -> tuple[mx.array, mx.array, mx.array]: ... @staticmethod def compute_attention( query: mx.array, key: mx.array, value: mx.array, batch_size: int, num_heads: int, head_dim: int, mask: mx.array | None = ..., ) -> mx.array: ... @staticmethod def convert_key_padding_mask_to_additive_mask( mask: mx.array | None, joint_seq_len: int, txt_seq_len: int ) -> mx.array | None: ... @staticmethod def apply_rope( xq: mx.array, xk: mx.array, freqs_cis: mx.array ) -> tuple[mx.array, mx.array]: ... @staticmethod def apply_rope_bshd( xq: mx.array, xk: mx.array, cos: mx.array, sin: mx.array ) -> tuple[mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/embed_nd.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class EmbedND(nn.Module): def __init__(self) -> None: ... def __call__(self, ids: mx.array) -> mx.array: ... @staticmethod def rope(pos: mx.array, dim: int, theta: float) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/feed_forward.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class FeedForward(nn.Module): def __init__(self, activation_function) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/guidance_embedder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class GuidanceEmbedder(nn.Module): def __init__(self) -> None: ... def __call__(self, sample: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/joint_attention.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn from typing import Any class JointAttention(nn.Module): num_heads: int head_dimension: int to_q: nn.Linear to_k: nn.Linear to_v: nn.Linear norm_q: nn.RMSNorm norm_k: nn.RMSNorm add_q_proj: nn.Linear add_k_proj: nn.Linear add_v_proj: nn.Linear norm_added_q: nn.RMSNorm norm_added_k: nn.RMSNorm to_out: list[Any] to_add_out: nn.Linear def __init__(self) -> None: ... def __call__( self, hidden_states: mx.array, encoder_hidden_states: mx.array, image_rotary_emb: mx.array, ) -> tuple[mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/joint_transformer_block.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn from typing import Any from mflux.models.flux.model.flux_transformer.joint_attention import JointAttention from mflux.models.flux.model.flux_transformer.ada_layer_norm_zero import ( AdaLayerNormZero, ) class JointTransformerBlock(nn.Module): attn: JointAttention norm1: AdaLayerNormZero norm1_context: AdaLayerNormZero norm2: nn.Module norm2_context: nn.Module ff: nn.Module ff_context: nn.Module def __init__(self, layer: Any) -> None: ... def __call__( self, hidden_states: mx.array, encoder_hidden_states: mx.array, text_embeddings: mx.array, rotary_embeddings: mx.array, ) -> tuple[mx.array, mx.array]: ... @staticmethod def apply_norm_and_feed_forward( hidden_states: mx.array, attn_output: mx.array, gate_mlp: mx.array, gate_msa: mx.array, scale_mlp: mx.array, shift_mlp: mx.array, norm_layer: nn.Module, ff_layer: nn.Module, ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/single_block_attention.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class SingleBlockAttention(nn.Module): num_heads: int head_dimension: int to_q: nn.Linear to_k: nn.Linear to_v: nn.Linear norm_q: nn.RMSNorm norm_k: nn.RMSNorm def __init__(self) -> None: ... def __call__( self, hidden_states: mx.array, image_rotary_emb: mx.array ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/single_transformer_block.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn from typing import Any from mflux.models.flux.model.flux_transformer.single_block_attention import ( SingleBlockAttention, ) from mflux.models.flux.model.flux_transformer.ada_layer_norm_zero_single import ( AdaLayerNormZeroSingle, ) class SingleTransformerBlock(nn.Module): attn: SingleBlockAttention norm: AdaLayerNormZeroSingle def __init__(self, layer: Any) -> None: ... def __call__( self, hidden_states: mx.array, text_embeddings: mx.array, rotary_embeddings: mx.array, ) -> tuple[mx.array, mx.array]: ... def _apply_feed_forward_and_projection( self, norm_hidden_states: mx.array, attn_output: mx.array, gate: mx.array ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/text_embedder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class TextEmbedder(nn.Module): def __init__(self) -> None: ... def __call__(self, caption: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/time_text_embed.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn from mflux.models.common.config import ModelConfig class TimeTextEmbed(nn.Module): def __init__(self, model_config: ModelConfig) -> None: ... def __call__( self, time_step: mx.array, pooled_projection: mx.array, guidance: mx.array ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/timestep_embedder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class TimestepEmbedder(nn.Module): def __init__(self) -> None: ... def __call__(self, sample: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_transformer/transformer.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn from mflux.models.common.config.config import Config from mflux.models.common.config.model_config import ModelConfig from mflux.models.flux.model.flux_transformer.embed_nd import EmbedND from mflux.models.flux.model.flux_transformer.time_text_embed import TimeTextEmbed from mflux.models.flux.model.flux_transformer.joint_transformer_block import ( JointTransformerBlock, ) from mflux.models.flux.model.flux_transformer.single_transformer_block import ( SingleTransformerBlock, ) class Transformer(nn.Module): transformer_blocks: list[JointTransformerBlock] single_transformer_blocks: list[SingleTransformerBlock] x_embedder: nn.Linear pos_embed: EmbedND time_text_embed: TimeTextEmbed norm_out: nn.LayerNorm proj_out: nn.Linear context_embedder: nn.Linear def __init__( self, model_config: ModelConfig, num_transformer_blocks: int = ..., num_single_transformer_blocks: int = ..., ) -> None: ... def __call__( self, t: int, config: Config, hidden_states: mx.array, prompt_embeds: mx.array, pooled_prompt_embeds: mx.array, controlnet_block_samples: list[mx.array] | None = ..., controlnet_single_block_samples: list[mx.array] | None = ..., kontext_image_ids: mx.array | None = ..., ) -> mx.array: ... @staticmethod def compute_rotary_embeddings( prompt_embeds: mx.array, pos_embed: EmbedND, config: Config, kontext_image_ids: mx.array | None = ..., ) -> mx.array: ... @staticmethod def compute_text_embeddings( t: int, pooled_prompt_embeds: mx.array, time_text_embed: TimeTextEmbed, config: Config, ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/common/attention.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class Attention(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/common/resnet_block_2d.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class ResnetBlock2D(nn.Module): def __init__( self, norm1: int, conv1_in: int, conv1_out: int, norm2: int, conv2_in: int, conv2_out: int, conv_shortcut_in: int | None = ..., conv_shortcut_out: int | None = ..., is_conv_shortcut: bool = ..., ) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/common/unet_mid_block.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class UnetMidBlock(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/decoder/conv_in.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class ConvIn(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/decoder/conv_norm_out.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class ConvNormOut(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/decoder/conv_out.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class ConvOut(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/decoder/decoder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class Decoder(nn.Module): def __init__( self, enable_tiling: bool = ..., split_direction: str = ... ) -> None: ... def __call__(self, latents: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/decoder/up_block_1_or_2.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class UpBlock1Or2(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/decoder/up_block_3.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class UpBlock3(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/decoder/up_block_4.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class UpBlock4(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/decoder/up_sampler.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class UpSampler(nn.Module): def __init__(self, conv_in: int, conv_out: int) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... @staticmethod def up_sample_nearest(x: mx.array, scale: int = ...): # -> array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/encoder/conv_in.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class ConvIn(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/encoder/conv_norm_out.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class ConvNormOut(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/encoder/conv_out.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import mlx.nn as nn class ConvOut(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/encoder/down_block_1.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class DownBlock1(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/encoder/down_block_2.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class DownBlock2(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/encoder/down_block_3.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class DownBlock3(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/encoder/down_block_4.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class DownBlock4(nn.Module): def __init__(self) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/encoder/down_sampler.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class DownSampler(nn.Module): def __init__(self, conv_in: int, conv_out: int) -> None: ... def __call__(self, input_array: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/encoder/encoder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class Encoder(nn.Module): def __init__(self) -> None: ... def __call__(self, latents: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/flux_vae/vae.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class VAE(nn.Module): scaling_factor: int = ... shift_factor: int = ... spatial_scale = ... latent_channels = ... def __init__(self) -> None: ... def decode(self, latents: mx.array) -> mx.array: ... def encode(self, image: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/redux_encoder/redux_encoder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class ReduxEncoder(nn.Module): def __init__(self) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/siglip_vision_transformer/siglip_encoder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class SiglipEncoder(nn.Module): def __init__(self) -> None: ... def __call__(self, inputs_embeds: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/siglip_vision_transformer/siglip_encoder_layer.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class SiglipEncoderLayer(nn.Module): def __init__(self) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/siglip_vision_transformer/siglip_mlp.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class SiglipMLP(nn.Module): def __init__(self) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/siglip_vision_transformer/siglip_multi_head_attention_pooling_head.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class SiglipMultiHeadAttentionPoolingHead(nn.Module): def __init__(self) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/siglip_vision_transformer/siglip_sdpa_attention.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class SiglipSdpaAttention(nn.Module): head_dimension = ... batch_size = ... num_heads = ... def __init__(self) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... @staticmethod def reshape_and_transpose(x, batch_size, num_heads, head_dim): # -> array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/siglip_vision_transformer/siglip_vision_embeddings.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class SiglipVisionEmbeddings(nn.Module): embed_dim = ... image_size = ... patch_size = ... def __init__(self) -> None: ... def __call__(self, pixel_values: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/model/siglip_vision_transformer/siglip_vision_transformer.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class SiglipVisionTransformer(nn.Module): def __init__(self) -> None: ... def __call__(self, pixel_values: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/flux/variants/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/flux/variants/concept_attention/attention_data.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import PIL.Image from dataclasses import dataclass from pathlib import Path from typing import List from mflux.models.flux.variants.concept_attention.joint_transformer_block_concept import ( LayerAttentionData, ) @dataclass class TimestepAttentionData: t: int attention_information: List[LayerAttentionData] def stack_img_attentions(self) -> mx.array: ... def stack_concept_attentions(self) -> mx.array: ... class GenerationAttentionData: def __init__(self) -> None: ... def append(self, timestep_attention: TimestepAttentionData): # -> None: ... def stack_all_img_attentions(self) -> mx.array: ... def stack_all_concept_attentions(self) -> mx.array: ... @dataclass class ConceptHeatmap: concept: str image: PIL.Image.Image layer_indices: List[int] timesteps: List[int] height: int width: int def save( self, path: str | Path, export_json_metadata: bool = ..., overwrite: bool = ... ) -> None: ... def get_metadata(self) -> dict: ... ================================================ FILE: .mlx_typings/mflux/models/flux/variants/concept_attention/joint_attention_concept.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class JointAttentionConcept(nn.Module): def __init__(self) -> None: ... def __call__( self, hidden_states: mx.array, encoder_hidden_states: mx.array, encoder_hidden_states_concept: mx.array, image_rotary_emb: mx.array, image_rotary_emb_concept: mx.array, ) -> tuple[mx.array, mx.array, mx.array, mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/flux/variants/concept_attention/joint_transformer_block_concept.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from dataclasses import dataclass from mlx import nn @dataclass class LayerAttentionData: layer: int img_attention: mx.array concept_attention: mx.array ... class JointTransformerBlockConcept(nn.Module): def __init__(self, layer) -> None: ... def __call__( self, layer_idx: int, hidden_states: mx.array, encoder_hidden_states: mx.array, encoder_hidden_states_concept: mx.array, text_embeddings: mx.array, text_embeddings_concept: mx.array, rotary_embeddings: mx.array, rotary_embeddings_concept: mx.array, ) -> tuple[mx.array, mx.array, mx.array, LayerAttentionData]: ... ================================================ FILE: .mlx_typings/mflux/models/flux/variants/concept_attention/transformer_concept.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn from mflux.models.common.config.config import Config from mflux.models.common.config.model_config import ModelConfig from mflux.models.flux.variants.concept_attention.attention_data import ( TimestepAttentionData, ) class TransformerConcept(nn.Module): def __init__( self, model_config: ModelConfig, num_transformer_blocks: int = ..., num_single_transformer_blocks: int = ..., ) -> None: ... def __call__( self, t: int, config: Config, hidden_states: mx.array, prompt_embeds: mx.array, prompt_embeds_concept: mx.array, pooled_prompt_embeds: mx.array, pooled_prompt_embeds_concept: mx.array, ) -> tuple[mx.array, TimestepAttentionData]: ... ================================================ FILE: .mlx_typings/mflux/models/flux/variants/controlnet/transformer_controlnet.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn from mflux.models.common.config.config import Config from mflux.models.common.config.model_config import ModelConfig class TransformerControlnet(nn.Module): def __init__( self, model_config: ModelConfig, num_transformer_blocks: int = ..., num_single_transformer_blocks: int = ..., ) -> None: ... def __call__( self, t: int, config: Config, hidden_states: mx.array, prompt_embeds: mx.array, pooled_prompt_embeds: mx.array, controlnet_condition: mx.array, ) -> tuple[list[mx.array], list[mx.array]]: ... ================================================ FILE: .mlx_typings/mflux/models/flux/variants/kontext/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ from mflux.models.flux.variants.kontext.flux_kontext import Flux1Kontext __all__ = ["Flux1Kontext"] ================================================ FILE: .mlx_typings/mflux/models/flux/variants/kontext/flux_kontext.pyi ================================================ """ This type stub file was generated by pyright. """ from pathlib import Path from typing import Any from mlx import nn from mflux.models.common.config.model_config import ModelConfig from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import ( CLIPEncoder, ) from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5Encoder from mflux.models.flux.model.flux_transformer.transformer import Transformer from mflux.models.flux.model.flux_vae.vae import VAE from mflux.utils.generated_image import GeneratedImage class Flux1Kontext(nn.Module): vae: VAE transformer: Transformer t5_text_encoder: T5Encoder clip_text_encoder: CLIPEncoder bits: int | None lora_paths: list[str] | None lora_scales: list[float] | None prompt_cache: dict[str, Any] tokenizers: dict[str, Any] def __init__( self, quantize: int | None = ..., model_path: str | None = ..., lora_paths: list[str] | None = ..., lora_scales: list[float] | None = ..., model_config: ModelConfig = ..., ) -> None: ... def generate_image( self, seed: int, prompt: str, num_inference_steps: int = ..., height: int = ..., width: int = ..., guidance: float = ..., image_path: Path | str | None = ..., image_strength: float | None = ..., scheduler: str = ..., ) -> GeneratedImage: ... ================================================ FILE: .mlx_typings/mflux/models/flux/variants/kontext/kontext_util.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mflux.models.flux.model.flux_vae.vae import VAE class KontextUtil: @staticmethod def create_image_conditioning_latents( vae: VAE, height: int, width: int, image_path: str, ) -> tuple[mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/flux/variants/txt2img/flux.pyi ================================================ """ This type stub file was generated by pyright. """ from pathlib import Path from mlx import nn from typing import Any from mflux.models.common.config.model_config import ModelConfig from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import ( CLIPEncoder, ) from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5Encoder from mflux.models.flux.model.flux_transformer.transformer import Transformer from mflux.models.flux.model.flux_vae.vae import VAE from mflux.utils.generated_image import GeneratedImage class Flux1(nn.Module): vae: VAE transformer: Transformer t5_text_encoder: T5Encoder clip_text_encoder: CLIPEncoder bits: int | None lora_paths: list[str] | None lora_scales: list[float] | None prompt_cache: dict[str, Any] tokenizers: dict[str, Any] def __init__( self, quantize: int | None = ..., model_path: str | None = ..., lora_paths: list[str] | None = ..., lora_scales: list[float] | None = ..., model_config: ModelConfig = ..., ) -> None: ... def generate_image( self, seed: int, prompt: str, num_inference_steps: int = ..., height: int = ..., width: int = ..., guidance: float = ..., image_path: Path | str | None = ..., image_strength: float | None = ..., scheduler: str = ..., negative_prompt: str | None = ..., ) -> GeneratedImage: ... @staticmethod def from_name(model_name: str, quantize: int | None = ...) -> Flux1: ... def save_model(self, base_path: str) -> None: ... def freeze(self, **kwargs): # -> None: ... ================================================ FILE: .mlx_typings/mflux/models/flux/weights/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ from mflux.models.flux.weights.flux_weight_definition import FluxWeightDefinition from mflux.models.flux.weights.flux_weight_mapping import FluxWeightMapping __all__ = ["FluxWeightDefinition", "FluxWeightMapping"] ================================================ FILE: .mlx_typings/mflux/models/flux/weights/flux_lora_mapping.pyi ================================================ """ This type stub file was generated by pyright. """ from mflux.models.common.lora.mapping.lora_mapping import LoRAMapping, LoRATarget class FluxLoRAMapping(LoRAMapping): @staticmethod def get_mapping() -> list[LoRATarget]: ... ================================================ FILE: .mlx_typings/mflux/models/flux/weights/flux_weight_definition.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import List from mflux.models.common.weights.loading.weight_definition import ( ComponentDefinition, TokenizerDefinition, ) """ This type stub file was generated by pyright. """ class FluxWeightDefinition: @staticmethod def get_components() -> List[ComponentDefinition]: ... @staticmethod def get_tokenizers() -> List[TokenizerDefinition]: ... @staticmethod def get_download_patterns() -> List[str]: ... @staticmethod def quantization_predicate(path: str, module) -> bool: ... class FluxControlnetWeightDefinition: @staticmethod def get_controlnet_component() -> ComponentDefinition: ... @staticmethod def get_components() -> List[ComponentDefinition]: ... @staticmethod def get_tokenizers() -> List[TokenizerDefinition]: ... @staticmethod def get_download_patterns() -> List[str]: ... @staticmethod def quantization_predicate(path: str, module) -> bool: ... class FluxReduxWeightDefinition: @staticmethod def get_components() -> List[ComponentDefinition]: ... @staticmethod def get_tokenizers() -> List[TokenizerDefinition]: ... @staticmethod def get_download_patterns() -> List[str]: ... @staticmethod def quantization_predicate(path: str, module) -> bool: ... ================================================ FILE: .mlx_typings/mflux/models/flux/weights/flux_weight_mapping.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import List from mflux.models.common.weights.mapping.weight_mapping import ( WeightMapping, WeightTarget, ) class FluxWeightMapping(WeightMapping): @staticmethod def get_transformer_mapping() -> List[WeightTarget]: ... @staticmethod def get_controlnet_transformer_mapping() -> List[WeightTarget]: ... @staticmethod def get_vae_mapping() -> List[WeightTarget]: ... @staticmethod def get_t5_encoder_mapping() -> List[WeightTarget]: ... @staticmethod def get_clip_encoder_mapping() -> List[WeightTarget]: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/qwen/cli/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/qwen/latent_creator/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/qwen/latent_creator/qwen_latent_creator.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx """ This type stub file was generated by pyright. """ class QwenLatentCreator: @staticmethod def create_noise(seed: int, height: int, width: int) -> mx.array: ... @staticmethod def pack_latents( latents: mx.array, height: int, width: int, num_channels_latents: int = ... ) -> mx.array: ... @staticmethod def unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_attention.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenAttention(nn.Module): def __init__( self, hidden_size: int, num_attention_heads: int, num_key_value_heads: int = ..., max_position_embeddings: int = ..., rope_theta: float = ..., rope_scaling: dict = ..., ) -> None: ... def __call__( self, hidden_states: mx.array, attention_mask: mx.array | None = ..., position_embeddings: tuple[mx.array, mx.array] | None = ..., ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_encoder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenEncoder(nn.Module): def __init__( self, vocab_size: int = ..., hidden_size: int = ..., num_hidden_layers: int = ..., max_position_embeddings: int = ..., rope_theta: float = ..., ) -> None: ... def get_image_features( self, pixel_values: mx.array, image_grid_thw: mx.array ) -> mx.array: ... def __call__( self, input_ids: mx.array, attention_mask: mx.array, pixel_values: mx.array | None = ..., image_grid_thw: mx.array | None = ..., ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_encoder_layer.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenEncoderLayer(nn.Module): def __init__( self, hidden_size: int = ..., num_attention_heads: int = ..., num_key_value_heads: int = ..., intermediate_size: int = ..., rms_norm_eps: float = ..., max_position_embeddings: int = ..., rope_theta: float = ..., ) -> None: ... def __call__( self, hidden_states: mx.array, attention_mask: mx.array | None = ..., position_embeddings: tuple[mx.array, mx.array] | None = ..., ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_mlp.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenMLP(nn.Module): def __init__(self, hidden_size: int, intermediate_size: int) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_patch_merger.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class PatchMerger(nn.Module): def __init__( self, context_dim: int, hidden_size: int, spatial_merge_size: int = ... ) -> None: ... def __call__(self, x: mx.array, grid_thw: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_prompt_encoder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mflux.models.common.tokenizer import Tokenizer from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenTextEncoder """ This type stub file was generated by pyright. """ class QwenPromptEncoder: @staticmethod def encode_prompt( prompt: str, negative_prompt: str, prompt_cache: dict[str, tuple[mx.array, mx.array, mx.array, mx.array]], qwen_tokenizer: Tokenizer, qwen_text_encoder: QwenTextEncoder, ) -> tuple[mx.array, mx.array, mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_rms_norm.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenRMSNorm(nn.Module): def __init__(self, hidden_size: int, eps: float = ...) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_rope.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenRotaryEmbedding(nn.Module): def __init__( self, dim: int, max_position_embeddings: int = ..., base: float = ..., device: str = ..., scaling_factor: float = ..., rope_type: str = ..., config=..., ) -> None: ... def __call__( self, x: mx.array, position_ids: mx.array ) -> tuple[mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_text_encoder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn """ This type stub file was generated by pyright. """ class QwenTextEncoder(nn.Module): def __init__(self) -> None: ... def __call__( self, input_ids: mx.array, attention_mask: mx.array ) -> tuple[mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_attention.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class VisionAttention(nn.Module): def __init__(self, embed_dim: int = ..., num_heads: int = ...) -> None: ... def __call__( self, x: mx.array, position_embeddings=..., cu_seqlens=... ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_block.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class VisionBlock(nn.Module): def __init__( self, embed_dim: int = ..., num_heads: int = ..., mlp_ratio: float = ... ) -> None: ... def __call__( self, x: mx.array, position_embeddings=..., cu_seqlens=... ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_language_encoder.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenVisionLanguageEncoder(nn.Module): def __init__(self, encoder=...) -> None: ... def __call__( self, input_ids: mx.array, attention_mask: mx.array | None = ..., pixel_values: mx.array | None = ..., image_grid_thw: mx.array | None = ..., ) -> tuple[mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_mlp.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class VisionMLP(nn.Module): def __init__(self, dim: int, hidden_dim: int) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_patch_embed.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class VisionPatchEmbed(nn.Module): def __init__( self, patch_size: int = ..., temporal_patch_size: int = ..., in_channels: int = ..., embed_dim: int = ..., ) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_rotary_embedding.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class VisionRotaryEmbedding(nn.Module): def __init__(self, dim: int, theta: float = ...) -> None: ... def __call__(self, max_grid_size: int) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_transformer.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class VisionTransformer(nn.Module): def __init__( self, patch_size: int = ..., temporal_patch_size: int = ..., in_channels: int = ..., embed_dim: int = ..., depth: int = ..., num_heads: int = ..., mlp_ratio: float = ..., hidden_size: int = ..., spatial_merge_size: int = ..., window_size: int = ..., fullatt_block_indexes: list = ..., ) -> None: ... def get_window_index(self, grid_thw: mx.array): # -> tuple[array, array]: ... def rot_pos_emb(self, grid_thw: mx.array) -> mx.array: ... def __call__(self, pixel_values: mx.array, grid_thw: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_transformer/qwen_attention.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn from typing import Any class QwenAttention(nn.Module): _num_heads: int _head_dim: int num_heads: int head_dim: int to_q: nn.Linear to_k: nn.Linear to_v: nn.Linear add_q_proj: nn.Linear add_k_proj: nn.Linear add_v_proj: nn.Linear norm_q: nn.RMSNorm norm_k: nn.RMSNorm norm_added_q: nn.RMSNorm norm_added_k: nn.RMSNorm attn_to_out: list[Any] to_add_out: nn.Linear def __init__( self, dim: int = ..., num_heads: int = ..., head_dim: int = ... ) -> None: ... def __call__( self, img_modulated: mx.array, txt_modulated: mx.array, encoder_hidden_states_mask: mx.array | None, image_rotary_emb: tuple[mx.array, mx.array], block_idx: int | None = ..., ) -> tuple[mx.array, mx.array]: ... def _compute_attention_qwen( self, query: mx.array, key: mx.array, value: mx.array, mask: mx.array | None, block_idx: int | None, ) -> mx.array: ... @staticmethod def _convert_mask_for_qwen( mask: mx.array | None, joint_seq_len: int, txt_seq_len: int ) -> mx.array | None: ... @staticmethod def _apply_rope_qwen( x: mx.array, cos_vals: mx.array, sin_vals: mx.array ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_transformer/qwen_feed_forward.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenFeedForward(nn.Module): def __init__(self, dim: int = ...) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_transformer/qwen_rope.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenEmbedRopeMLX(nn.Module): def __init__( self, theta: int, axes_dim: list[int], scale_rope: bool = ... ) -> None: ... def __call__( self, video_fhw: tuple[int, int, int] | list[tuple[int, int, int]], txt_seq_lens: list[int], ) -> tuple[tuple[mx.array, mx.array], tuple[mx.array, mx.array]]: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_transformer/qwen_time_text_embed.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenTimeTextEmbed(nn.Module): def __init__(self, timestep_proj_dim: int = ..., inner_dim: int = ...) -> None: ... def __call__(self, timestep: mx.array, hidden_states: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_transformer/qwen_timestep_embedding.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenTimestepEmbedding(nn.Module): def __init__(self, proj_dim: int, inner_dim: int) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_transformer/qwen_timesteps.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenTimesteps(nn.Module): def __init__(self, proj_dim: int = ..., scale: float = ...) -> None: ... def __call__(self, timesteps: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_transformer/qwen_transformer.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn from typing import Any from mflux.models.common.config.config import Config from mflux.models.qwen.model.qwen_transformer.qwen_transformer_block import ( QwenTransformerBlock, ) class QwenTransformer(nn.Module): transformer_blocks: list[QwenTransformerBlock] inner_dim: int img_in: nn.Linear txt_in: nn.Linear txt_norm: nn.RMSNorm time_text_embed: Any pos_embed: Any norm_out: nn.Module proj_out: nn.Linear def __init__( self, in_channels: int = ..., out_channels: int = ..., num_layers: int = ..., attention_head_dim: int = ..., num_attention_heads: int = ..., joint_attention_dim: int = ..., patch_size: int = ..., ) -> None: ... def __call__( self, t: int, config: Config, hidden_states: mx.array, encoder_hidden_states: mx.array, encoder_hidden_states_mask: mx.array, qwen_image_ids: mx.array | None = ..., cond_image_grid: tuple[int, int, int] | None = ..., ) -> mx.array: ... @staticmethod def _compute_timestep(t: int | float, config: Config) -> mx.array: ... @staticmethod def _compute_rotary_embeddings( encoder_hidden_states_mask: mx.array, pos_embed: Any, config: Config, cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]] | None = ..., ) -> tuple[mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_transformer/qwen_transformer_block.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn from typing import Any from mflux.models.qwen.model.qwen_transformer.qwen_attention import QwenAttention class QwenTransformerBlock(nn.Module): attn: QwenAttention img_mod_linear: nn.Linear img_mod_silu: nn.SiLU txt_mod_linear: nn.Linear txt_mod_silu: nn.SiLU img_norm1: nn.RMSNorm txt_norm1: nn.RMSNorm img_norm2: nn.RMSNorm txt_norm2: nn.RMSNorm img_ff: Any txt_ff: Any def __init__( self, dim: int = ..., num_heads: int = ..., head_dim: int = ... ) -> None: ... def __call__( self, hidden_states: mx.array, encoder_hidden_states: mx.array, encoder_hidden_states_mask: mx.array | None, text_embeddings: mx.array, image_rotary_emb: tuple[mx.array, mx.array], block_idx: int | None = ..., ) -> tuple[mx.array, mx.array]: ... @staticmethod def _modulate(x: mx.array, mod_params: mx.array) -> tuple[mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_transformer/qwen_transformer_rms_norm.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenTransformerRMSNorm(nn.Module): def __init__(self, dim: int, eps: float = ...) -> None: ... def __call__(self, hidden_states: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_vae/qwen_image_attention_block_3d.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenImageAttentionBlock3D(nn.Module): def __init__(self, dim: int) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_vae/qwen_image_causal_conv_3d.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenImageCausalConv3D(nn.Module): def __init__( self, in_channels: int, out_channels: int, kernel_size: int = ..., stride: int = ..., padding: int = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_vae/qwen_image_decoder_3d.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenImageDecoder3D(nn.Module): def __init__(self) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_vae/qwen_image_down_block_3d.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenImageDownBlock3D(nn.Module): def __init__( self, in_channels: int, out_channels: int, num_res_blocks: int = ..., downsample_mode: str = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_vae/qwen_image_encoder_3d.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenImageEncoder3D(nn.Module): def __init__(self) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_vae/qwen_image_mid_block_3d.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenImageMidBlock3D(nn.Module): def __init__(self, dim: int, num_layers: int = ...) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_vae/qwen_image_res_block_3d.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenImageResBlock3D(nn.Module): def __init__(self, in_channels: int, out_channels: int) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_vae/qwen_image_resample_3d.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenImageResample3D(nn.Module): def __init__(self, dim: int, mode: str) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_vae/qwen_image_rms_norm.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenImageRMSNorm(nn.Module): def __init__( self, num_channels: int, eps: float = ..., images: bool = ... ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_vae/qwen_image_up_block_3d.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenImageUpBlock3D(nn.Module): def __init__( self, in_channels: int, out_channels: int, num_res_blocks: int = ..., upsample_mode: str = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/model/qwen_vae/qwen_vae.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mlx import nn class QwenVAE(nn.Module): LATENTS_MEAN = ... LATENTS_STD = ... spatial_scale = ... latent_channels = ... def __init__(self) -> None: ... def decode(self, latents: mx.array) -> mx.array: ... def encode(self, latents: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/qwen_initializer.pyi ================================================ """ This type stub file was generated by pyright. """ from mflux.models.common.config import ModelConfig class QwenImageInitializer: @staticmethod def init( model, model_config: ModelConfig, quantize: int | None, model_path: str | None = ..., lora_paths: list[str] | None = ..., lora_scales: list[float] | None = ..., ) -> None: ... @staticmethod def init_edit( model, model_config: ModelConfig, quantize: int | None, model_path: str | None = ..., lora_paths: list[str] | None = ..., lora_scales: list[float] | None = ..., ) -> None: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/tokenizer/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/qwen/tokenizer/qwen_image_processor.pyi ================================================ """ This type stub file was generated by pyright. """ import numpy as np from typing import Optional, Union from PIL import Image OPENAI_CLIP_MEAN = ... OPENAI_CLIP_STD = ... def smart_resize( height: int, width: int, factor: int = ..., min_pixels: int = ..., max_pixels: int = ..., ) -> tuple[int, int]: ... class QwenImageProcessor: def __init__( self, min_pixels: int = ..., max_pixels: int = ..., patch_size: int = ..., temporal_patch_size: int = ..., merge_size: int = ..., image_mean: Optional[list[float]] = ..., image_std: Optional[list[float]] = ..., ) -> None: ... def preprocess( self, images: Union[Image.Image, list[Image.Image]] ) -> tuple[np.ndarray, np.ndarray]: ... def get_number_of_image_patches( self, height: int, width: int, min_pixels: Optional[int] = ..., max_pixels: Optional[int] = ..., ) -> int: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/tokenizer/qwen_vision_language_processor.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Optional, Union from PIL import Image from mflux.models.qwen.tokenizer.qwen_image_processor import QwenImageProcessor class QwenVisionLanguageProcessor: def __init__( self, tokenizer, image_processor: Optional[QwenImageProcessor] = ..., image_token: str = ..., video_token: str = ..., ) -> None: ... def __call__( self, images: Optional[Union[Image.Image, list[Image.Image]]] = ..., text: Optional[Union[str, list[str]]] = ..., padding: bool = ..., return_tensors: Optional[str] = ..., ) -> dict: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/tokenizer/qwen_vision_language_tokenizer.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import numpy as np from typing import Union from PIL import Image from mflux.models.qwen.tokenizer.qwen_vision_language_processor import ( QwenVisionLanguageProcessor, ) class QwenVisionLanguageTokenizer: def __init__( self, processor: QwenVisionLanguageProcessor, max_length: int = ..., use_picture_prefix: bool = ..., ) -> None: ... def tokenize_with_image( self, prompt: str, image: Union[Image.Image, np.ndarray, str, list], vl_width: int | None = ..., vl_height: int | None = ..., ) -> tuple[mx.array, mx.array, mx.array, mx.array]: ... def tokenize_text_only(self, prompt: str) -> tuple[mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/variants/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/models/qwen/variants/edit/qwen_edit_util.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from mflux.models.common.vae.tiling_config import TilingConfig class QwenEditUtil: @staticmethod def create_image_conditioning_latents( vae, height: int, width: int, image_paths: list[str] | str, vl_width: int | None = ..., vl_height: int | None = ..., tiling_config: TilingConfig | None = ..., ) -> tuple[mx.array, mx.array, int, int, int]: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/variants/edit/qwen_image_edit.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from pathlib import Path from mlx import nn from typing import Any from mflux.models.common.config import Config from mflux.models.common.config.model_config import ModelConfig from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenTextEncoder from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer from mflux.models.qwen.model.qwen_vae.qwen_vae import QwenVAE from mflux.utils.generated_image import GeneratedImage class QwenImageEdit(nn.Module): vae: QwenVAE transformer: QwenTransformer text_encoder: QwenTextEncoder bits: int | None lora_paths: list[str] | None lora_scales: list[float] | None prompt_cache: dict[str, Any] tokenizers: dict[str, Any] def __init__( self, quantize: int | None = ..., model_path: str | None = ..., lora_paths: list[str] | None = ..., lora_scales: list[float] | None = ..., model_config: ModelConfig = ..., ) -> None: ... def generate_image( self, seed: int, prompt: str, image_paths: list[str], num_inference_steps: int = ..., height: int | None = ..., width: int | None = ..., guidance: float = ..., image_path: Path | str | None = ..., scheduler: str = ..., negative_prompt: str | None = ..., ) -> GeneratedImage: ... def _encode_prompts_with_images( self, prompt: str, negative_prompt: str, image_paths: list[str], config: Config, vl_width: int | None, vl_height: int | None, ) -> tuple[mx.array, mx.array, mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/variants/txt2img/qwen_image.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from pathlib import Path from mlx import nn from typing import Any from mflux.models.common.config import ModelConfig from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenTextEncoder from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer from mflux.models.qwen.model.qwen_vae.qwen_vae import QwenVAE from mflux.utils.generated_image import GeneratedImage class QwenImage(nn.Module): vae: QwenVAE transformer: QwenTransformer text_encoder: QwenTextEncoder bits: int | None lora_paths: list[str] | None lora_scales: list[float] | None prompt_cache: dict[str, Any] tokenizers: dict[str, Any] def __init__( self, quantize: int | None = ..., model_path: str | None = ..., lora_paths: list[str] | None = ..., lora_scales: list[float] | None = ..., model_config: ModelConfig = ..., ) -> None: ... def generate_image( self, seed: int, prompt: str, num_inference_steps: int = ..., height: int = ..., width: int = ..., guidance: float = ..., image_path: Path | str | None = ..., image_strength: float | None = ..., scheduler: str = ..., negative_prompt: str | None = ..., ) -> GeneratedImage: ... def save_model(self, base_path: str) -> None: ... @staticmethod def compute_guided_noise( noise: mx.array, noise_negative: mx.array, guidance: float ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/weights/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ from mflux.models.qwen.weights.qwen_weight_definition import QwenWeightDefinition from mflux.models.qwen.weights.qwen_weight_mapping import QwenWeightMapping __all__ = ["QwenWeightDefinition", "QwenWeightMapping"] ================================================ FILE: .mlx_typings/mflux/models/qwen/weights/qwen_lora_mapping.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import List from mflux.models.common.lora.mapping.lora_mapping import LoRAMapping, LoRATarget class QwenLoRAMapping(LoRAMapping): @staticmethod def get_mapping() -> List[LoRATarget]: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/weights/qwen_weight_definition.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import List from mflux.models.common.weights.loading.weight_definition import ( ComponentDefinition, TokenizerDefinition, ) """ This type stub file was generated by pyright. """ class QwenWeightDefinition: @staticmethod def get_components() -> List[ComponentDefinition]: ... @staticmethod def get_tokenizers() -> List[TokenizerDefinition]: ... @staticmethod def get_download_patterns() -> List[str]: ... @staticmethod def quantization_predicate(path: str, module) -> bool: ... ================================================ FILE: .mlx_typings/mflux/models/qwen/weights/qwen_weight_mapping.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import List from mflux.models.common.weights.mapping.weight_mapping import ( WeightMapping, WeightTarget, ) class QwenWeightMapping(WeightMapping): @staticmethod def get_transformer_mapping() -> List[WeightTarget]: ... @staticmethod def get_vae_mapping() -> List[WeightTarget]: ... @staticmethod def get_text_encoder_mapping() -> List[WeightTarget]: ... @staticmethod def get_mapping() -> List[WeightTarget]: ... ================================================ FILE: .mlx_typings/mflux/models/seedvr2/weights/seedvr2_weight_definition.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import List from mflux.models.common.weights.loading.weight_definition import ( ComponentDefinition, TokenizerDefinition, ) """ This type stub file was generated by pyright. """ class SeedVR2WeightDefinition: @staticmethod def get_components() -> List[ComponentDefinition]: ... @staticmethod def get_tokenizers() -> List[TokenizerDefinition]: ... @staticmethod def get_download_patterns() -> List[str]: ... @staticmethod def quantization_predicate(path: str, module) -> bool: ... ================================================ FILE: .mlx_typings/mflux/models/seedvr2/weights/seedvr2_weight_mapping.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import List from mflux.models.common.weights.mapping.weight_mapping import ( WeightMapping, WeightTarget, ) class SeedVR2WeightMapping(WeightMapping): @staticmethod def get_transformer_mapping() -> List[WeightTarget]: ... @staticmethod def get_vae_mapping() -> List[WeightTarget]: ... ================================================ FILE: .mlx_typings/mflux/models/z_image/latent_creator/z_image_latent_creator.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx class ZImageLatentCreator: @staticmethod def create_noise(seed: int, height: int, width: int) -> mx.array: ... @staticmethod def pack_latents(latents: mx.array, height: int, width: int) -> mx.array: ... @staticmethod def unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/models/z_image/weights/z_image_weight_definition.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import List from mflux.models.common.weights.loading.weight_definition import ( ComponentDefinition, TokenizerDefinition, ) """ This type stub file was generated by pyright. """ class ZImageWeightDefinition: @staticmethod def get_components() -> List[ComponentDefinition]: ... @staticmethod def get_tokenizers() -> List[TokenizerDefinition]: ... @staticmethod def get_download_patterns() -> List[str]: ... @staticmethod def quantization_predicate(path: str, module) -> bool: ... ================================================ FILE: .mlx_typings/mflux/models/z_image/weights/z_image_weight_mapping.pyi ================================================ """ This type stub file was generated by pyright. """ from mflux.models.common.weights.mapping.weight_mapping import ( WeightMapping, WeightTarget, ) class ZImageWeightMapping(WeightMapping): @staticmethod def get_text_encoder_mapping() -> list[WeightTarget]: ... @staticmethod def get_vae_mapping() -> list[WeightTarget]: ... @staticmethod def get_transformer_mapping() -> list[WeightTarget]: ... ================================================ FILE: .mlx_typings/mflux/release/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/utils/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ """ This type stub file was generated by pyright. """ ================================================ FILE: .mlx_typings/mflux/utils/box_values.pyi ================================================ """ This type stub file was generated by pyright. """ from dataclasses import dataclass @dataclass class AbsoluteBoxValues: top: int right: int bottom: int left: int ... class BoxValueError(ValueError): ... @dataclass class BoxValues: top: int | str right: int | str bottom: int | str left: int | str def normalize_to_dimensions(self, width, height) -> AbsoluteBoxValues: ... @staticmethod def parse(value, delimiter=...) -> BoxValues: ... ================================================ FILE: .mlx_typings/mflux/utils/exceptions.pyi ================================================ """ This type stub file was generated by pyright. """ """ This type stub file was generated by pyright. """ class MFluxException(Exception): ... class ImageSavingException(MFluxException): ... class MetadataEmbedException(MFluxException): ... class MFluxUserException(MFluxException): ... class PromptFileReadError(MFluxUserException): ... class StopImageGenerationException(MFluxUserException): ... class StopTrainingException(MFluxUserException): ... class CommandExecutionError(MFluxException): def __init__( self, cmd: list[str], return_code: int, stdout: str | None, stderr: str | None ) -> None: ... class ReferenceVsOutputImageError(AssertionError): ... class ModelConfigError(ValueError): ... class InvalidBaseModel(ModelConfigError): ... ================================================ FILE: .mlx_typings/mflux/utils/generated_image.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import PIL.Image from pathlib import Path from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.concept_attention.attention_data import ConceptHeatmap log = ... class GeneratedImage: image: PIL.Image.Image def __init__( self, image: PIL.Image.Image, model_config: ModelConfig, seed: int, prompt: str, steps: int, guidance: float | None, precision: mx.Dtype, quantization: int, generation_time: float, lora_paths: list[str] | None = ..., lora_scales: list[float] | None = ..., height: int | None = ..., width: int | None = ..., controlnet_image_path: str | Path | None = ..., controlnet_strength: float | None = ..., image_path: str | Path | None = ..., image_paths: list[str] | list[Path] | None = ..., image_strength: float | None = ..., masked_image_path: str | Path | None = ..., depth_image_path: str | Path | None = ..., redux_image_paths: list[str] | list[Path] | None = ..., redux_image_strengths: list[float] | None = ..., concept_heatmap: ConceptHeatmap | None = ..., negative_prompt: str | None = ..., init_metadata: dict | None = ..., ) -> None: ... def get_right_half(self) -> GeneratedImage: ... def save( self, path: str | Path, export_json_metadata: bool = ..., overwrite: bool = ... ) -> None: ... def save_with_heatmap( self, path: str | Path, export_json_metadata: bool = ..., overwrite: bool = ... ) -> None: ... def save_concept_heatmap( self, path: str | Path, export_json_metadata: bool = ..., overwrite: bool = ... ) -> None: ... ================================================ FILE: .mlx_typings/mflux/utils/image_util.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx import PIL.Image from pathlib import Path from typing import Any from PIL._typing import StrOrBytesPath from mflux.models.common.config.config import Config from mflux.models.flux.variants.concept_attention.attention_data import ConceptHeatmap from mflux.utils.box_values import AbsoluteBoxValues from mflux.utils.generated_image import GeneratedImage log = ... class ImageUtil: @staticmethod def to_image( decoded_latents: mx.array, config: Config, seed: int, prompt: str, quantization: int, generation_time: float, lora_paths: list[str] | None = ..., lora_scales: list[float] | None = ..., controlnet_image_path: str | Path | None = ..., image_path: str | Path | None = ..., image_paths: list[str] | list[Path] | None = ..., redux_image_paths: list[str] | list[Path] | None = ..., redux_image_strengths: list[float] | None = ..., image_strength: float | None = ..., masked_image_path: str | Path | None = ..., depth_image_path: str | Path | None = ..., concept_heatmap: ConceptHeatmap | None = ..., negative_prompt: str | None = ..., init_metadata: dict[str, Any] | None = ..., ) -> GeneratedImage: ... @staticmethod def to_composite_image( generated_images: list[GeneratedImage], ) -> PIL.Image.Image: ... @staticmethod def to_array(image: PIL.Image.Image, is_mask: bool = ...) -> mx.array: ... @staticmethod def load_image( image_or_path: PIL.Image.Image | StrOrBytesPath, ) -> PIL.Image.Image: ... @staticmethod def expand_image( image: PIL.Image.Image, box_values: AbsoluteBoxValues | None = ..., top: int | str = ..., right: int | str = ..., bottom: int | str = ..., left: int | str = ..., fill_color: tuple = ..., ) -> PIL.Image.Image: ... @staticmethod def create_outpaint_mask_image( orig_width: int, orig_height: int, **create_bordered_image_kwargs ): # -> Image: ... @staticmethod def create_bordered_image( orig_width: int, orig_height: int, border_color: tuple, content_color: tuple, box_values: AbsoluteBoxValues | None = ..., top: int | str = ..., right: int | str = ..., bottom: int | str = ..., left: int | str = ..., ) -> PIL.Image.Image: ... @staticmethod def scale_to_dimensions( image: PIL.Image.Image, target_width: int, target_height: int ) -> PIL.Image.Image: ... @staticmethod def save_image( image: PIL.Image.Image, path: str | Path, metadata: dict | None = ..., export_json_metadata: bool = ..., overwrite: bool = ..., ) -> None: ... @staticmethod def preprocess_for_model( image: PIL.Image.Image, target_size: tuple = ..., mean: list = ..., std: list = ..., resample: int = ..., ) -> mx.array: ... @staticmethod def preprocess_for_depth_pro( image: PIL.Image.Image, target_size: tuple = ..., mean: list = ..., std: list = ..., resample: int = ..., ) -> mx.array: ... ================================================ FILE: .mlx_typings/mflux/utils/metadata_builder.pyi ================================================ """ This type stub file was generated by pyright. """ from pathlib import Path log = ... class MetadataBuilder: _IPTC_PROMPT_MAX_BYTES = ... @staticmethod def embed_metadata(metadata: dict, path: str | Path) -> None: ... @staticmethod def build_xmp_packet(metadata: dict) -> str: ... @staticmethod def build_iptc_binary(metadata: dict) -> bytes: ... ================================================ FILE: .mlx_typings/mflux/utils/version_util.pyi ================================================ """ This type stub file was generated by pyright. """ """ This type stub file was generated by pyright. """ class VersionUtil: @staticmethod def get_mflux_version() -> str: ... ================================================ FILE: .mlx_typings/mlx/core/__init__.pyi ================================================ import enum import pathlib import types from typing import ( Annotated, Callable, Literal, Mapping, Sequence, TypeAlias, overload, ) import numpy from mlx.nn.layers import Module from numpy.typing import ArrayLike as _ArrayLike from . import cuda as cuda from . import distributed as distributed from . import metal as metal from . import random as random class ArrayAt: """A helper object to apply updates at specific indices.""" def __getitem__(self, indices: object | None) -> ArrayAt: ... def add( self, value: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def subtract( self, value: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def multiply( self, value: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def divide( self, value: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def maximum( self, value: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def minimum( self, value: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... class ArrayIterator: """A helper object to iterate over the 1st dimension of an array.""" def __next__(self) -> array: ... def __iter__(self) -> ArrayIterator: ... class ArrayLike: """ Any Python object which has an ``__mlx__array__`` method that returns an :obj:`array`. """ def __init__(self, arg: object, /) -> None: ... class Device: """A device to run operations on.""" def __init__(self, type: DeviceType, index: int = ...) -> None: ... @property def type(self) -> DeviceType: ... def __repr__(self) -> str: ... def __eq__(self, arg: object, /) -> bool: ... class DeviceType(enum.Enum): cpu = ... # type: ignore gpu = ... #  type: ignore def __eq__(self, arg: object, /) -> bool: ... class Dtype: """ An object to hold the type of a :class:`array`. See the :ref:`list of types ` for more details on available data types. """ @property def size(self) -> int: """Size of the type in bytes.""" def __repr__(self) -> str: ... def __eq__(self, arg: object, /) -> bool: ... def __hash__(self) -> int: ... class DtypeCategory(enum.Enum): """ Type to hold categories of :class:`dtypes `. * :attr:`~mlx.core.generic` * :ref:`bool_ ` * :attr:`~mlx.core.number` * :attr:`~mlx.core.integer` * :attr:`~mlx.core.unsignedinteger` * :ref:`uint8 ` * :ref:`uint16 ` * :ref:`uint32 ` * :ref:`uint64 ` * :attr:`~mlx.core.signedinteger` * :ref:`int8 ` * :ref:`int32 ` * :ref:`int64 ` * :attr:`~mlx.core.inexact` * :attr:`~mlx.core.floating` * :ref:`float16 ` * :ref:`bfloat16 ` * :ref:`float32 ` * :ref:`float64 ` * :attr:`~mlx.core.complexfloating` * :ref:`complex64 ` See also :func:`~mlx.core.issubdtype`. """ complexfloating = ... floating = ... inexact = ... signedinteger = ... unsignedinteger = ... integer = ... number = ... generic = ... class FunctionExporter: """ A context managing class for exporting multiple traces of the same function to a file. Make an instance of this class by calling fun:`mx.exporter`. """ def close(self) -> None: ... def __enter__(self) -> FunctionExporter: ... def __exit__( self, exc_type: object | None = ..., exc_value: object | None = ..., traceback: object | None = ..., ) -> None: ... def __call__(self, *args, **kwargs) -> None: ... class Stream: """A stream for running operations on a given device.""" @property def device(self) -> Device: ... def __repr__(self) -> str: ... def __eq__(self, arg: object, /) -> bool: ... class StreamContext: """ A context manager for setting the current device and stream. See :func:`stream` for usage. Args: s: The stream or device to set as the default. """ def __init__(self, s: Stream | Device) -> None: ... def __enter__(self) -> None: ... def __exit__( self, exc_type: type | None = ..., exc_value: object | None = ..., traceback: object | None = ..., ) -> None: ... def device_info() -> dict[str, str | int]: """ Get information about the GPU device and system settings. Currently returns: * ``architecture`` * ``max_buffer_size`` * ``max_recommended_working_set_size`` * ``memory_size`` * ``resource_limit`` Returns: dict: A dictionary with string keys and string or integer values. """ def abs(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise absolute value. Args: a (array): Input array. Returns: array: The absolute value of ``a``. """ def add( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise addition. Add two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The sum of ``a`` and ``b``. """ def addmm( c: array, a: array, b: array, /, alpha: float = ..., beta: float = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Matrix multiplication with addition and optional scaling. Perform the (possibly batched) matrix multiplication of two arrays and add to the result with optional scaling factors. Args: c (array): Input array or scalar. a (array): Input array or scalar. b (array): Input array or scalar. alpha (float, optional): Scaling factor for the matrix product of ``a`` and ``b`` (default: ``1``) beta (float, optional): Scaling factor for ``c`` (default: ``1``) Returns: array: ``alpha * (a @ b) + beta * c`` """ def all( a: array, /, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """ An `and` reduction over the given axes. Args: a (array): Input array. axis (int or list(int), optional): Optional axis or axes to reduce over. If unspecified this defaults to reducing over the entire array. keepdims (bool, optional): Keep reduced axes as singleton dimensions, defaults to `False`. Returns: array: The output array with the corresponding axes reduced. """ def allclose( a: array, b: array, /, rtol: float = ..., atol: float = ..., *, equal_nan: bool = ..., stream: Stream | Device | None = ..., ) -> array: """ Approximate comparison of two arrays. Infinite values are considered equal if they have the same sign, NaN values are not equal unless ``equal_nan`` is ``True``. The arrays are considered equal if: .. code-block:: all(abs(a - b) <= (atol + rtol * abs(b))) Note unlike :func:`array_equal`, this function supports numpy-style broadcasting. Args: a (array): Input array. b (array): Input array. rtol (float): Relative tolerance. atol (float): Absolute tolerance. equal_nan (bool): If ``True``, NaNs are considered equal. Defaults to ``False``. Returns: array: The boolean output scalar indicating if the arrays are close. """ def any( a: array, /, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """ An `or` reduction over the given axes. Args: a (array): Input array. axis (int or list(int), optional): Optional axis or axes to reduce over. If unspecified this defaults to reducing over the entire array. keepdims (bool, optional): Keep reduced axes as singleton dimensions, defaults to `False`. Returns: array: The output array with the corresponding axes reduced. """ @overload def arange( start: int | float, stop: int | float, step: int | float | None, dtype: Dtype | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Generates ranges of numbers. Generate numbers in the half-open interval ``[start, stop)`` in increments of ``step``. Args: start (float or int, optional): Starting value which defaults to ``0``. stop (float or int): Stopping value. step (float or int, optional): Increment which defaults to ``1``. dtype (Dtype, optional): Specifies the data type of the output. If unspecified will default to ``float32`` if any of ``start``, ``stop``, or ``step`` are ``float``. Otherwise will default to ``int32``. Returns: array: The range of values. Note: Following the Numpy convention the actual increment used to generate numbers is ``dtype(start + step) - dtype(start)``. This can lead to unexpected results for example if `start + step` is a fractional value and the `dtype` is integral. """ @overload def arange( stop: int | float, step: int | float | None = ..., dtype: Dtype | None = ..., *, stream: Stream | Device | None = ..., ) -> array: ... def arccos(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise inverse cosine. Args: a (array): Input array. Returns: array: The inverse cosine of ``a``. """ def arccosh(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise inverse hyperbolic cosine. Args: a (array): Input array. Returns: array: The inverse hyperbolic cosine of ``a``. """ def arcsin(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise inverse sine. Args: a (array): Input array. Returns: array: The inverse sine of ``a``. """ def arcsinh(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise inverse hyperbolic sine. Args: a (array): Input array. Returns: array: The inverse hyperbolic sine of ``a``. """ def arctan(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise inverse tangent. Args: a (array): Input array. Returns: array: The inverse tangent of ``a``. """ def arctan2(a: array, b: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise inverse tangent of the ratio of two arrays. Args: a (array): Input array. b (array): Input array. Returns: array: The inverse tangent of the ratio of ``a`` and ``b``. """ def arctanh(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise inverse hyperbolic tangent. Args: a (array): Input array. Returns: array: The inverse hyperbolic tangent of ``a``. """ def argmax( a: array, /, axis: int | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Indices of the maximum values along the axis. Args: a (array): Input array. axis (int, optional): Optional axis to reduce over. If unspecified this defaults to reducing over the entire array. keepdims (bool, optional): Keep reduced axes as singleton dimensions, defaults to `False`. Returns: array: The ``uint32`` array with the indices of the maximum values. """ def argmin( a: array, /, axis: int | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Indices of the minimum values along the axis. Args: a (array): Input array. axis (int, optional): Optional axis to reduce over. If unspecified this defaults to reducing over the entire array. keepdims (bool, optional): Keep reduced axes as singleton dimensions, defaults to `False`. Returns: array: The ``uint32`` array with the indices of the minimum values. """ def argpartition( a: array, /, kth: int, axis: int | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Returns the indices that partition the array. The ordering of the elements within a partition in given by the indices is undefined. Args: a (array): Input array. kth (int): Element index at the ``kth`` position in the output will give the sorted position. All indices before the ``kth`` position will be of elements less or equal to the element at the ``kth`` index and all indices after will be of elements greater or equal to the element at the ``kth`` index. axis (int or None, optional): Optional axis to partition over. If ``None``, this partitions over the flattened array. If unspecified, it defaults to ``-1``. Returns: array: The ``uint32`` array containing indices that partition the input. """ def argsort( a: array, /, axis: int | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Returns the indices that sort the array. Args: a (array): Input array. axis (int or None, optional): Optional axis to sort over. If ``None``, this sorts over the flattened array. If unspecified, it defaults to -1 (sorting over the last axis). Returns: array: The ``uint32`` array containing indices that sort the input. """ class array: """An N-dimensional array object.""" def __init__( self: array, val: scalar | list | tuple | numpy.ndarray | array, dtype: Dtype | None = ..., ) -> None: ... def __buffer__(self, flags, /): """ Return a buffer object that exposes the underlying memory of the object. """ def __release_buffer__(self, buffer, /): """ Release the buffer object that exposes the underlying memory of the object. """ @property def size(self) -> int: """Number of elements in the array.""" @property def ndim(self) -> int: """The array's dimension.""" @property def itemsize(self) -> int: """The size of the array's datatype in bytes.""" @property def nbytes(self) -> int: """The number of bytes in the array.""" @property def shape(self) -> tuple[int, ...]: """ The shape of the array as a Python tuple. Returns: tuple(int): A tuple containing the sizes of each dimension. """ @property def dtype(self) -> Dtype: """The array's :class:`Dtype`.""" @property def real(self) -> array: """The real part of a complex array.""" @property def imag(self) -> array: """The imaginary part of a complex array.""" def item(self) -> scalar: """ Access the value of a scalar array. Returns: Standard Python scalar. """ def tolist(self) -> list_or_scalar: """ Convert the array to a Python :class:`list`. Returns: list: The Python list. If the array is a scalar then a standard Python scalar is returned. If the array has more than one dimension then the result is a nested list of lists. The value type of the list corresponding to the last dimension is either ``bool``, ``int`` or ``float`` depending on the ``dtype`` of the array. """ def astype(self, dtype: Dtype, stream: Stream | Device | None = ...) -> array: """ Cast the array to a specified type. Args: dtype (Dtype): Type to which the array is cast. stream (Stream): Stream (or device) for the operation. Returns: array: The array with type ``dtype``. """ def __array_namespace__(self, api_version: str | None = ...) -> types.ModuleType: """ Returns an object that has all the array API functions on it. See the `Python array API `_ for more information. Args: api_version (str, optional): String representing the version of the array API spec to return. Default: ``None``. Returns: out (Any): An object representing the array API namespace. """ def __getitem__(self, arg: object | None) -> array: ... def __setitem__( self, arg0: object | None, arg1: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> None: ... @property def at(self) -> ArrayAt: """ Used to apply updates at the given indices. .. note:: Regular in-place updates map to assignment. For instance ``x[idx] += y`` maps to ``x[idx] = x[idx] + y``. As a result, assigning to the same index ignores all but one update. Using ``x.at[idx].add(y)`` will correctly apply all updates to all indices. .. list-table:: :header-rows: 1 * - array.at syntax - In-place syntax * - ``x = x.at[idx].add(y)`` - ``x[idx] += y`` * - ``x = x.at[idx].subtract(y)`` - ``x[idx] -= y`` * - ``x = x.at[idx].multiply(y)`` - ``x[idx] *= y`` * - ``x = x.at[idx].divide(y)`` - ``x[idx] /= y`` * - ``x = x.at[idx].maximum(y)`` - ``x[idx] = mx.maximum(x[idx], y)`` * - ``x = x.at[idx].minimum(y)`` - ``x[idx] = mx.minimum(x[idx], y)`` Example: >>> a = mx.array([0, 0]) >>> idx = mx.array([0, 1, 0, 1]) >>> a[idx] += 1 >>> a array([1, 1], dtype=int32) >>> >>> a = mx.array([0, 0]) >>> a.at[idx].add(1) array([2, 2], dtype=int32) """ def __len__(self) -> int: ... def __iter__(self) -> ArrayIterator: ... def __getstate__(self) -> tuple: ... def __setstate__(self, arg: tuple, /) -> None: ... def __dlpack__(self) -> _ArrayLike: ... def __dlpack_device__(self) -> tuple: ... def __copy__(self) -> array: ... def __deepcopy__(self, memo: dict) -> array: ... def __add__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __iadd__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __radd__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __sub__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __isub__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __rsub__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __mul__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __imul__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __rmul__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __truediv__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __itruediv__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __rtruediv__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __div__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __rdiv__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __floordiv__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __ifloordiv__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __rfloordiv__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __mod__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __imod__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __rmod__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __eq__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array | bool: ... def __lt__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __le__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __gt__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __ge__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __ne__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array | bool: ... def __neg__(self) -> array: ... def __bool__(self) -> bool: ... def __repr__(self) -> str: ... def __matmul__(self, other: array) -> array: ... def __imatmul__(self, other: array) -> array: ... def __pow__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __rpow__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __ipow__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __invert__(self) -> array: ... def __and__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __iand__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __or__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __ior__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __lshift__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __ilshift__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __rshift__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __irshift__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __xor__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __ixor__( self, other: bool | int | float | array | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] | complex | ArrayLike, ) -> array: ... def __int__(self) -> int: ... def __float__(self) -> float: ... def flatten( self, start_axis: int = ..., end_axis: int = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`flatten`.""" def reshape(self, *shape: int, stream: Stream | Device | None = ...) -> array: """ Equivalent to :func:`reshape` but the shape can be passed either as a :obj:`tuple` or as separate arguments. See :func:`reshape` for full documentation. """ def squeeze( self, axis: int | Sequence[int] | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`squeeze`.""" def abs(self, *, stream: Stream | Device | None = ...) -> array: """See :func:`abs`.""" def __abs__(self) -> array: """See :func:`abs`.""" def square(self, *, stream: Stream | Device | None = ...) -> array: """See :func:`square`.""" def sqrt(self, *, stream: Stream | Device | None = ...) -> array: """See :func:`sqrt`.""" def rsqrt(self, *, stream: Stream | Device | None = ...) -> array: """See :func:`rsqrt`.""" def reciprocal(self, *, stream: Stream | Device | None = ...) -> array: """See :func:`reciprocal`.""" def exp(self, *, stream: Stream | Device | None = ...) -> array: """See :func:`exp`.""" def log(self, *, stream: Stream | Device | None = ...) -> array: """See :func:`log`.""" def log2(self, *, stream: Stream | Device | None = ...) -> array: """See :func:`log2`.""" def log10(self, *, stream: Stream | Device | None = ...) -> array: """See :func:`log10`.""" def sin(self, *, stream: Stream | Device | None = ...) -> array: """See :func:`sin`.""" def cos(self, *, stream: Stream | Device | None = ...) -> array: """See :func:`cos`.""" def log1p(self, *, stream: Stream | Device | None = ...) -> array: """See :func:`log1p`.""" def all( self, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`all`.""" def any( self, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`any`.""" def moveaxis( self, source: int, destination: int, *, stream: Stream | Device | None = ... ) -> array: """See :func:`moveaxis`.""" def swapaxes( self, axis1: int, axis2: int, *, stream: Stream | Device | None = ... ) -> array: """See :func:`swapaxes`.""" def transpose(self, *axes: int, stream: Stream | Device | None = ...) -> array: """ Equivalent to :func:`transpose` but the axes can be passed either as a tuple or as separate arguments. See :func:`transpose` for full documentation. """ @property def T(self) -> array: """Equivalent to calling ``self.transpose()`` with no arguments.""" def sum( self, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`sum`.""" def prod( self, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`prod`.""" def min( self, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`min`.""" def max( self, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`max`.""" def logcumsumexp( self, axis: int | None = ..., *, reverse: bool = ..., inclusive: bool = ..., stream: Stream | Device | None = ..., ) -> array: """See :func:`logcumsumexp`.""" def logsumexp( self, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`logsumexp`.""" def mean( self, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`mean`.""" def std( self, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., ddof: int = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`std`.""" def var( self, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., ddof: int = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`var`.""" def split( self, indices_or_sections: int | tuple[int, ...], axis: int = ..., *, stream: Stream | Device | None = ..., ) -> list[array]: """See :func:`split`.""" def argmin( self, axis: int | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`argmin`.""" def argmax( self, axis: int | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`argmax`.""" def cumsum( self, axis: int | None = ..., *, reverse: bool = ..., inclusive: bool = ..., stream: Stream | Device | None = ..., ) -> array: """See :func:`cumsum`.""" def cumprod( self, axis: int | None = ..., *, reverse: bool = ..., inclusive: bool = ..., stream: Stream | Device | None = ..., ) -> array: """See :func:`cumprod`.""" def cummax( self, axis: int | None = ..., *, reverse: bool = ..., inclusive: bool = ..., stream: Stream | Device | None = ..., ) -> array: """See :func:`cummax`.""" def cummin( self, axis: int | None = ..., *, reverse: bool = ..., inclusive: bool = ..., stream: Stream | Device | None = ..., ) -> array: """See :func:`cummin`.""" def round( self, decimals: int = ..., *, stream: Stream | Device | None = ... ) -> array: """See :func:`round`.""" def diagonal( self, offset: int = ..., axis1: int = ..., axis2: int = ..., stream: Stream | Device | None = ..., ) -> array: """See :func:`diagonal`.""" def diag(self, k: int = ..., *, stream: Stream | Device | None = ...) -> array: """Extract a diagonal or construct a diagonal matrix.""" def conj(self, *, stream: Stream | Device | None = ...) -> array: """See :func:`conj`.""" def view(self, dtype: Dtype, *, stream: Stream | Device | None = ...) -> array: """See :func:`view`.""" def array_equal( a: scalar | array, b: scalar | array, equal_nan: bool = ..., stream: Stream | Device | None = ..., ) -> array: """ Array equality check. Compare two arrays for equality. Returns ``True`` if and only if the arrays have the same shape and their values are equal. The arrays need not have the same type to be considered equal. Args: a (array): Input array or scalar. b (array): Input array or scalar. equal_nan (bool): If ``True``, NaNs are considered equal. Defaults to ``False``. Returns: array: A scalar boolean array. """ def as_strided( a: array, /, shape: Sequence[int] | None = ..., strides: Sequence[int] | None = ..., offset: int = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Create a view into the array with the given shape and strides. The resulting array will always be as if the provided array was row contiguous regardless of the provided arrays storage order and current strides. .. note:: Note that this function should be used with caution as it changes the shape and strides of the array directly. This can lead to the resulting array pointing to invalid memory locations which can result into crashes. Args: a (array): Input array shape (list(int), optional): The shape of the resulting array. If None it defaults to ``a.shape()``. strides (list(int), optional): The strides of the resulting array. If None it defaults to the reverse exclusive cumulative product of ``a.shape()``. offset (int): Skip that many elements from the beginning of the input array. Returns: array: The output array which is the strided view of the input. """ def async_eval(*args: MX_ARRAY_TREE) -> None: """ Asynchronously evaluate an :class:`array` or tree of :class:`array`. .. note:: This is an experimental API and may change in future versions. Args: *args (arrays or trees of arrays): Each argument can be a single array or a tree of arrays. If a tree is given the nodes can be a Python :class:`list`, :class:`tuple` or :class:`dict`. Leaves which are not arrays are ignored. Example: >>> x = mx.array(1.0) >>> y = mx.exp(x) >>> mx.async_eval(y) >>> print(y) >>> >>> y = mx.exp(x) >>> mx.async_eval(y) >>> z = y + 3 >>> mx.async_eval(z) >>> print(z) """ def atleast_1d( *arys: array, stream: Stream | Device | None = ... ) -> array | list[array]: """ Convert all arrays to have at least one dimension. Args: *arys: Input arrays. stream (Stream | Device | None, optional): The stream to execute the operation on. Returns: array or list(array): An array or list of arrays with at least one dimension. """ def atleast_2d( *arys: array, stream: Stream | Device | None = ... ) -> array | list[array]: """ Convert all arrays to have at least two dimensions. Args: *arys: Input arrays. stream (Stream | Device | None, optional): The stream to execute the operation on. Returns: array or list(array): An array or list of arrays with at least two dimensions. """ def atleast_3d( *arys: array, stream: Stream | Device | None = ... ) -> array | list[array]: """ Convert all arrays to have at least three dimensions. Args: *arys: Input arrays. stream (Stream | Device | None, optional): The stream to execute the operation on. Returns: array or list(array): An array or list of arrays with at least three dimensions. """ bfloat16: Dtype = ... def bitwise_and( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise bitwise and. Take the bitwise and of two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The bitwise and ``a & b``. """ def bitwise_invert(a: scalar | array, stream: Stream | Device | None = ...) -> array: """ Element-wise bitwise inverse. Take the bitwise complement of the input. Args: a (array): Input array or scalar. Returns: array: The bitwise inverse ``~a``. """ def bitwise_or( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise bitwise or. Take the bitwise or of two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The bitwise or``a | b``. """ def bitwise_xor( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise bitwise xor. Take the bitwise exclusive or of two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The bitwise xor ``a ^ b``. """ def block_masked_mm( a: array, b: array, /, block_size: int = ..., mask_out: array | None = ..., mask_lhs: array | None = ..., mask_rhs: array | None = ..., *, stream: Stream | Device | None = ..., ) -> array: r""" Matrix multiplication with block masking. Perform the (possibly batched) matrix multiplication of two arrays and with blocks of size ``block_size x block_size`` optionally masked out. Assuming ``a`` with shape (..., `M`, `K`) and b with shape (..., `K`, `N`) * ``lhs_mask`` must have shape (..., :math:`\lceil` `M` / ``block_size`` :math:`\rceil`, :math:`\lceil` `K` / ``block_size`` :math:`\rceil`) * ``rhs_mask`` must have shape (..., :math:`\lceil` `K` / ``block_size`` :math:`\rceil`, :math:`\lceil` `N` / ``block_size`` :math:`\rceil`) * ``out_mask`` must have shape (..., :math:`\lceil` `M` / ``block_size`` :math:`\rceil`, :math:`\lceil` `N` / ``block_size`` :math:`\rceil`) Note: Only ``block_size=64`` and ``block_size=32`` are currently supported Args: a (array): Input array or scalar. b (array): Input array or scalar. block_size (int): Size of blocks to be masked. Must be ``32`` or ``64``. Default: ``64``. mask_out (array, optional): Mask for output. Default: ``None``. mask_lhs (array, optional): Mask for ``a``. Default: ``None``. mask_rhs (array, optional): Mask for ``b``. Default: ``None``. Returns: array: The output array. """ def broadcast_arrays( *arrays: array, stream: Stream | Device | None = ... ) -> tuple[array, ...]: """ Broadcast arrays against one another. The broadcasting semantics are the same as Numpy. Args: *arrays (array): The input arrays. Returns: tuple(array): The output arrays with the broadcasted shape. """ def broadcast_shapes(*shapes: Sequence[int]) -> tuple[int]: """ Broadcast shapes. Returns the shape that results from broadcasting the supplied array shapes against each other. Args: *shapes (Sequence[int]): The shapes to broadcast. Returns: tuple: The broadcasted shape. Raises: ValueError: If the shapes cannot be broadcast. Example: >>> mx.broadcast_shapes((1,), (3, 1)) (3, 1) >>> mx.broadcast_shapes((6, 7), (5, 6, 1), (7,)) (5, 6, 7) >>> mx.broadcast_shapes((5, 1, 4), (1, 3, 1)) (5, 3, 4) """ def broadcast_to( a: scalar | array, /, shape: Sequence[int], *, stream: Stream | Device | None = ..., ) -> array: """ Broadcast an array to the given shape. The broadcasting semantics are the same as Numpy. Args: a (array): Input array. shape (list(int)): The shape to broadcast to. Returns: array: The output array with the new shape. """ def ceil(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise ceil. Args: a (array): Input array. Returns: array: The ceil of ``a``. """ def checkpoint(fun: Callable) -> Callable: ... def clear_cache() -> None: """ Clear the memory cache. After calling this, :func:`get_cache_memory` should return ``0``. """ def clip( a: array, /, a_min: scalar | array | None, a_max: scalar | array | None, *, stream: Stream | Device | None = ..., ) -> array: """ Clip the values of the array between the given minimum and maximum. If either ``a_min`` or ``a_max`` are ``None``, then corresponding edge is ignored. At least one of ``a_min`` and ``a_max`` cannot be ``None``. The input ``a`` and the limits must broadcast with one another. Args: a (array): Input array. a_min (scalar or array or None): Minimum value to clip to. a_max (scalar or array or None): Maximum value to clip to. Returns: array: The clipped array. """ def compile( fun: Callable, inputs: object | None = ..., outputs: object | None = ..., shapeless: bool = ..., ) -> Callable: """ Returns a compiled function which produces the same output as ``fun``. Args: fun (Callable): A function which takes a variable number of :class:`array` or trees of :class:`array` and returns a variable number of :class:`array` or trees of :class:`array`. inputs (list or dict, optional): These inputs will be captured during the function compilation along with the inputs to ``fun``. The ``inputs`` can be a :obj:`list` or a :obj:`dict` containing arbitrarily nested lists, dictionaries, or arrays. Leaf nodes that are not :obj:`array` are ignored. Default: ``None`` outputs (list or dict, optional): These outputs will be captured and updated in a compiled function. The ``outputs`` can be a :obj:`list` or a :obj:`dict` containing arbitrarily nested lists, dictionaries, or arrays. Leaf nodes that are not :obj:`array` are ignored. Default: ``None`` shapeless (bool, optional): A function compiled with the ``shapeless`` option enabled will not be recompiled when the input shape changes. Not all functions can be compiled with ``shapeless`` enabled. Attempting to compile such functions with shapeless enabled will throw. Note, changing the number of dimensions or type of any input will result in a recompilation even with ``shapeless`` set to ``True``. Default: ``False`` Returns: Callable: A compiled function which has the same input arguments as ``fun`` and returns the the same output(s). """ complex64: Dtype = ... complexfloating: DtypeCategory = ... def concat( arrays: list[array], axis: int | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`concatenate`.""" def concatenate( arrays: list[array], axis: int | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Concatenate the arrays along the given axis. Args: arrays (list(array)): Input :obj:`list` or :obj:`tuple` of arrays. axis (int, optional): Optional axis to concatenate along. If unspecified defaults to ``0``. Returns: array: The concatenated array. """ def conj(a: array, *, stream: Stream | Device | None = ...) -> array: """ Return the elementwise complex conjugate of the input. Alias for `mx.conjugate`. Args: a (array): Input array Returns: array: The output array. """ def conjugate(a: array, *, stream: Stream | Device | None = ...) -> array: """ Return the elementwise complex conjugate of the input. Alias for `mx.conj`. Args: a (array): Input array Returns: array: The output array. """ def contiguous( a: array, /, allow_col_major: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Force an array to be row contiguous. Copy if necessary. Args: a (array): The input to make contiguous allow_col_major (bool): Consider column major as contiguous and don't copy Returns: array: The row or col contiguous output. """ def conv1d( input: array, weight: array, /, stride: int = ..., padding: int = ..., dilation: int = ..., groups: int = ..., *, stream: Stream | Device | None = ..., ) -> array: """ 1D convolution over an input with several channels Args: input (array): Input array of shape ``(N, L, C_in)``. weight (array): Weight array of shape ``(C_out, K, C_in)``. stride (int, optional): Kernel stride. Default: ``1``. padding (int, optional): Input padding. Default: ``0``. dilation (int, optional): Kernel dilation. Default: ``1``. groups (int, optional): Input feature groups. Default: ``1``. Returns: array: The convolved array. """ def conv2d( input: array, weight: array, /, stride: int | tuple[int, int] = ..., padding: int | tuple[int, int] = ..., dilation: int | tuple[int, int] = ..., groups: int = ..., *, stream: Stream | Device | None = ..., ) -> array: """ 2D convolution over an input with several channels Args: input (array): Input array of shape ``(N, H, W, C_in)``. weight (array): Weight array of shape ``(C_out, KH, KW, C_in)``. stride (int or tuple(int), optional): :obj:`tuple` of size 2 with kernel strides. All spatial dimensions get the same stride if only one number is specified. Default: ``1``. padding (int or tuple(int), optional): :obj:`tuple` of size 2 with symmetric input padding. All spatial dimensions get the same padding if only one number is specified. Default: ``0``. dilation (int or tuple(int), optional): :obj:`tuple` of size 2 with kernel dilation. All spatial dimensions get the same dilation if only one number is specified. Default: ``1`` groups (int, optional): input feature groups. Default: ``1``. Returns: array: The convolved array. """ def conv3d( input: array, weight: array, /, stride: int | tuple[int, int, int] = ..., padding: int | tuple[int, int, int] = ..., dilation: int | tuple[int, int, int] = ..., groups: int = ..., *, stream: Stream | Device | None = ..., ) -> array: """ 3D convolution over an input with several channels Note: Only the default ``groups=1`` is currently supported. Args: input (array): Input array of shape ``(N, D, H, W, C_in)``. weight (array): Weight array of shape ``(C_out, KD, KH, KW, C_in)``. stride (int or tuple(int), optional): :obj:`tuple` of size 3 with kernel strides. All spatial dimensions get the same stride if only one number is specified. Default: ``1``. padding (int or tuple(int), optional): :obj:`tuple` of size 3 with symmetric input padding. All spatial dimensions get the same padding if only one number is specified. Default: ``0``. dilation (int or tuple(int), optional): :obj:`tuple` of size 3 with kernel dilation. All spatial dimensions get the same dilation if only one number is specified. Default: ``1`` groups (int, optional): input feature groups. Default: ``1``. Returns: array: The convolved array. """ def conv_general( input: array, weight: array, /, stride: int | Sequence[int] = ..., padding: int | Sequence[int] | tuple[Sequence[int] | Sequence[int]] = ..., kernel_dilation: int | Sequence[int] = ..., input_dilation: int | Sequence[int] = ..., groups: int = ..., flip: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """ General convolution over an input with several channels Args: input (array): Input array of shape ``(N, ..., C_in)``. weight (array): Weight array of shape ``(C_out, ..., C_in)``. stride (int or list(int), optional): :obj:`list` with kernel strides. All spatial dimensions get the same stride if only one number is specified. Default: ``1``. padding (int, list(int), or tuple(list(int), list(int)), optional): :obj:`list` with input padding. All spatial dimensions get the same padding if only one number is specified. Default: ``0``. kernel_dilation (int or list(int), optional): :obj:`list` with kernel dilation. All spatial dimensions get the same dilation if only one number is specified. Default: ``1`` input_dilation (int or list(int), optional): :obj:`list` with input dilation. All spatial dimensions get the same dilation if only one number is specified. Default: ``1`` groups (int, optional): Input feature groups. Default: ``1``. flip (bool, optional): Flip the order in which the spatial dimensions of the weights are processed. Performs the cross-correlation operator when ``flip`` is ``False`` and the convolution operator otherwise. Default: ``False``. Returns: array: The convolved array. """ def conv_transpose1d( input: array, weight: array, /, stride: int = ..., padding: int = ..., dilation: int = ..., output_padding: int = ..., groups: int = ..., *, stream: Stream | Device | None = ..., ) -> array: """ 1D transposed convolution over an input with several channels Args: input (array): Input array of shape ``(N, L, C_in)``. weight (array): Weight array of shape ``(C_out, K, C_in)``. stride (int, optional): Kernel stride. Default: ``1``. padding (int, optional): Input padding. Default: ``0``. dilation (int, optional): Kernel dilation. Default: ``1``. output_padding (int, optional): Output padding. Default: ``0``. groups (int, optional): Input feature groups. Default: ``1``. Returns: array: The convolved array. """ def conv_transpose2d( input: array, weight: array, /, stride: int | tuple[int, int] = ..., padding: int | tuple[int, int] = ..., dilation: int | tuple[int, int] = ..., output_padding: int | tuple[int, int] = ..., groups: int = ..., *, stream: Stream | Device | None = ..., ) -> array: """ 2D transposed convolution over an input with several channels Note: Only the default ``groups=1`` is currently supported. Args: input (array): Input array of shape ``(N, H, W, C_in)``. weight (array): Weight array of shape ``(C_out, KH, KW, C_in)``. stride (int or tuple(int), optional): :obj:`tuple` of size 2 with kernel strides. All spatial dimensions get the same stride if only one number is specified. Default: ``1``. padding (int or tuple(int), optional): :obj:`tuple` of size 2 with symmetric input padding. All spatial dimensions get the same padding if only one number is specified. Default: ``0``. dilation (int or tuple(int), optional): :obj:`tuple` of size 2 with kernel dilation. All spatial dimensions get the same dilation if only one number is specified. Default: ``1`` output_padding (int or tuple(int), optional): :obj:`tuple` of size 2 with output padding. All spatial dimensions get the same output padding if only one number is specified. Default: ``0``. groups (int, optional): input feature groups. Default: ``1``. Returns: array: The convolved array. """ def conv_transpose3d( input: array, weight: array, /, stride: int | tuple[int, int, int] = ..., padding: int | tuple[int, int, int] = ..., dilation: int | tuple[int, int, int] = ..., output_padding: int | tuple[int, int, int] = ..., groups: int = ..., *, stream: Stream | Device | None = ..., ) -> array: """ 3D transposed convolution over an input with several channels Note: Only the default ``groups=1`` is currently supported. Args: input (array): Input array of shape ``(N, D, H, W, C_in)``. weight (array): Weight array of shape ``(C_out, KD, KH, KW, C_in)``. stride (int or tuple(int), optional): :obj:`tuple` of size 3 with kernel strides. All spatial dimensions get the same stride if only one number is specified. Default: ``1``. padding (int or tuple(int), optional): :obj:`tuple` of size 3 with symmetric input padding. All spatial dimensions get the same padding if only one number is specified. Default: ``0``. dilation (int or tuple(int), optional): :obj:`tuple` of size 3 with kernel dilation. All spatial dimensions get the same dilation if only one number is specified. Default: ``1`` output_padding (int or tuple(int), optional): :obj:`tuple` of size 3 with output padding. All spatial dimensions get the same output padding if only one number is specified. Default: ``0``. groups (int, optional): input feature groups. Default: ``1``. Returns: array: The convolved array. """ def convolve( a: array, v: array, /, mode: str = ..., *, stream: Stream | Device | None = ... ) -> array: """ The discrete convolution of 1D arrays. If ``v`` is longer than ``a``, then they are swapped. The conv filter is flipped following signal processing convention. Args: a (array): 1D Input array. v (array): 1D Input array. mode (str, optional): {'full', 'valid', 'same'} Returns: array: The convolved array. """ def cos(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise cosine. Args: a (array): Input array. Returns: array: The cosine of ``a``. """ def cosh(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise hyperbolic cosine. Args: a (array): Input array. Returns: array: The hyperbolic cosine of ``a``. """ cpu: DeviceType = ... def cummax( a: array, /, axis: int | None = ..., *, reverse: bool = ..., inclusive: bool = ..., stream: Stream | Device | None = ..., ) -> array: """ Return the cumulative maximum of the elements along the given axis. Args: a (array): Input array axis (int, optional): Optional axis to compute the cumulative maximum over. If unspecified the cumulative maximum of the flattened array is returned. reverse (bool): Perform the cumulative maximum in reverse. inclusive (bool): The i-th element of the output includes the i-th element of the input. Returns: array: The output array. """ def cummin( a: array, /, axis: int | None = ..., *, reverse: bool = ..., inclusive: bool = ..., stream: Stream | Device | None = ..., ) -> array: """ Return the cumulative minimum of the elements along the given axis. Args: a (array): Input array axis (int, optional): Optional axis to compute the cumulative minimum over. If unspecified the cumulative minimum of the flattened array is returned. reverse (bool): Perform the cumulative minimum in reverse. inclusive (bool): The i-th element of the output includes the i-th element of the input. Returns: array: The output array. """ def cumprod( a: array, /, axis: int | None = ..., *, reverse: bool = ..., inclusive: bool = ..., stream: Stream | Device | None = ..., ) -> array: """ Return the cumulative product of the elements along the given axis. Args: a (array): Input array axis (int, optional): Optional axis to compute the cumulative product over. If unspecified the cumulative product of the flattened array is returned. reverse (bool): Perform the cumulative product in reverse. inclusive (bool): The i-th element of the output includes the i-th element of the input. Returns: array: The output array. """ def cumsum( a: array, /, axis: int | None = ..., *, reverse: bool = ..., inclusive: bool = ..., stream: Stream | Device | None = ..., ) -> array: """ Return the cumulative sum of the elements along the given axis. Args: a (array): Input array axis (int, optional): Optional axis to compute the cumulative sum over. If unspecified the cumulative sum of the flattened array is returned. reverse (bool): Perform the cumulative sum in reverse. inclusive (bool): The i-th element of the output includes the i-th element of the input. Returns: array: The output array. """ class custom_function: """ Set up a function for custom gradient and vmap definitions. This class is meant to be used as a function decorator. Instances are callables that behave identically to the wrapped function. However, when a function transformation is used (e.g. computing gradients using :func:`value_and_grad`) then the functions defined via :meth:`custom_function.vjp`, :meth:`custom_function.jvp` and :meth:`custom_function.vmap` are used instead of the default transformation. Note, all custom transformations are optional. Undefined transformations fall back to the default behaviour. Example: .. code-block:: python import mlx.core as mx @mx.custom_function def f(x, y): return mx.sin(x) * y @f.vjp def f_vjp(primals, cotangent, output): x, y = primals return cotan * mx.cos(x) * y, cotan * mx.sin(x) @f.jvp def f_jvp(primals, tangents): x, y = primals dx, dy = tangents return dx * mx.cos(x) * y + dy * mx.sin(x) @f.vmap def f_vmap(inputs, axes): x, y = inputs ax, ay = axes if ay != ax and ax is not None: y = y.swapaxes(ay, ax) return mx.sin(x) * y, (ax or ay) All ``custom_function`` instances behave as pure functions. Namely, any variables captured will be treated as constants and no gradients will be computed with respect to the captured arrays. For instance: .. code-block:: python import mlx.core as mx def g(x, y): @mx.custom_function def f(x): return x * y @f.vjp def f_vjp(x, dx, fx): # Note that we have only x, dx and fx and nothing with respect to y raise ValueError("Abort!") return f(x) x = mx.array(2.0) y = mx.array(3.0) print(g(x, y)) # prints 6.0 print(mx.grad(g)(x, y)) # Raises exception print(mx.grad(g, argnums=1)(x, y)) # prints 0.0 """ def __init__(self, f: Callable) -> None: ... def __call__(self, *args, **kwargs) -> object: ... def vjp(self, f: Callable): """ Define a custom vjp for the wrapped function. The vjp function takes three arguments: - *primals*: A pytree that contains all the positional arguments to the function. It could be a single array, a tuple of arrays or a full blown tuple of dicts of arrays etc. - *cotangents*: A pytree that matches the structure of the output but contains the cotangents (usually the gradients of the loss function with respect to the outputs). - *outputs*: The outputs of the function to be used to avoid recomputing them for the gradient computation. The vjp function should return the same pytree structure as the primals but containing the corresponding computed cotangents. """ def jvp(self, f: Callable): """ Define a custom jvp for the wrapped function. The jvp function takes two arguments: - *primals*: A pytree that contains all the positional arguments to the function. It could be a single array, a tuple of arrays or a full blown tuple of dicts of arrays etc. - *tangents*: A pytree that matches the structure of the inputs but instead contains the gradients wrt to each input. Tangents could be ``None`` if some inputs don't have an associated gradient. The jvp function should return the same pytree structure as the outputs of the function but containing the tangents. """ def vmap(self, f: Callable): """ Define a custom vectorization transformation for the wrapped function. The vmap function takes two arguments: - *inputs*: A pytree that contains all the positional arguments to the function. It could be a single array, a tuple of arrays or a full blown tuple of dicts of arrays etc. - *axes*: A pytree that matches the structure of the inputs but instead contains the vectorization axis for each input or ``None`` if an input is not vectorized. The vmap function should return the outputs of the original function but vectorized over the provided axes. It should also return a pytree with the vectorization axes of each output. If some outputs are no longer vectorized, then their vectorization axis should be ``None``. """ def default_device() -> Device: """Get the default device.""" def default_stream(device: Device | DeviceType) -> Stream: """Get the device's default stream.""" def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Convert angles from radians to degrees. Args: a (array): Input array. Returns: array: The angles in degrees. """ def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array]): """ Insert dependencies between arrays in the graph. The outputs are identical to ``inputs`` but with dependencies on ``dependencies``. Args: inputs (array or Sequence[array]): The input array or arrays. dependencies (array or Sequence[array]): The array or arrays to insert dependencies on. Returns: array or Sequence[array]: The outputs which depend on dependencies. """ def dequantize( w: array, /, scales: array, biases: array | None = ..., group_size: int = ..., bits: int = ..., mode: str = ..., *, stream: Stream | Device | None = ..., ) -> array: r""" Dequantize the matrix ``w`` using quantization parameters. Args: w (array): Matrix to be dequantized scales (array): The scales to use per ``group_size`` elements of ``w``. biases (array, optional): The biases to use per ``group_size`` elements of ``w``. Default: ``None``. group_size (int, optional): The size of the group in ``w`` that shares a scale and bias. Default: ``64``. bits (int, optional): The number of bits occupied by each element in ``w``. Default: ``4``. mode (str, optional): The quantization mode. Default: ``"affine"``. Returns: array: The dequantized version of ``w`` Notes: The currently supported quantization modes are ``"affine"`` and ``mxfp4``. For ``affine`` quantization, given the notation in :func:`quantize`, we compute :math:`w_i` from :math:`\hat{w_i}` and corresponding :math:`s` and :math:`\beta` as follows .. math:: w_i = s \hat{w_i} + \beta """ def diag(a: array, /, k: int = ..., *, stream: Stream | Device | None = ...) -> array: """ Extract a diagonal or construct a diagonal matrix. If ``a`` is 1-D then a diagonal matrix is constructed with ``a`` on the :math:`k`-th diagonal. If ``a`` is 2-D then the :math:`k`-th diagonal is returned. Args: a (array): 1-D or 2-D input array. k (int, optional): The diagonal to extract or construct. Default: ``0``. Returns: array: The extracted diagonal or the constructed diagonal matrix. """ def diagonal( a: array, offset: int = ..., axis1: int = ..., axis2: int = ..., stream: Stream | Device | None = ..., ) -> array: """ Return specified diagonals. If ``a`` is 2-D, then a 1-D array containing the diagonal at the given ``offset`` is returned. If ``a`` has more than two dimensions, then ``axis1`` and ``axis2`` determine the 2D subarrays from which diagonals are extracted. The new shape is the original shape with ``axis1`` and ``axis2`` removed and a new dimension inserted at the end corresponding to the diagonal. Args: a (array): Input array offset (int, optional): Offset of the diagonal from the main diagonal. Can be positive or negative. Default: ``0``. axis1 (int, optional): The first axis of the 2-D sub-arrays from which the diagonals should be taken. Default: ``0``. axis2 (int, optional): The second axis of the 2-D sub-arrays from which the diagonals should be taken. Default: ``1``. Returns: array: The diagonals of the array. """ def disable_compile() -> None: """ Globally disable compilation. Setting the environment variable ``MLX_DISABLE_COMPILE`` can also be used to disable compilation. """ def divide( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise division. Divide two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The quotient ``a / b``. """ def divmod( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise quotient and remainder. The fuction ``divmod(a, b)`` is equivalent to but faster than ``(a // b, a % b)``. The function uses numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: tuple(array, array): The quotient ``a // b`` and remainder ``a % b``. """ e: float = ... def einsum(subscripts: str, *operands, stream: Stream | Device | None = ...) -> array: """ Perform the Einstein summation convention on the operands. Args: subscripts (str): The Einstein summation convention equation. *operands (array): The input arrays. Returns: array: The output array. """ def einsum_path(subscripts: str, *operands): """ Compute the contraction order for the given Einstein summation. Args: subscripts (str): The Einstein summation convention equation. *operands (array): The input arrays. Returns: tuple(list(tuple(int, int)), str): The einsum path and a string containing information about the chosen path. """ def enable_compile() -> None: """ Globally enable compilation. This will override the environment variable ``MLX_DISABLE_COMPILE`` if set. """ def equal( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise equality. Equality comparison on two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The element-wise comparison ``a == b``. """ def erf(a: array, /, *, stream: Stream | Device | None = ...) -> array: r""" Element-wise error function. .. math:: \mathrm{erf}(x) = \frac{2}{\sqrt{\pi}} \int_0^x e^{-t^2} \, dt Args: a (array): Input array. Returns: array: The error function of ``a``. """ def erfinv(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise inverse of :func:`erf`. Args: a (array): Input array. Returns: array: The inverse error function of ``a``. """ euler_gamma: float = ... type MX_ARRAY_TREE = ( array | Module | list[MX_ARRAY_TREE] | tuple[MX_ARRAY_TREE, ...] | Mapping[str, MX_ARRAY_TREE] ) def eval(*args: MX_ARRAY_TREE | None) -> None: """ Evaluate an :class:`array` or tree of :class:`array`. Args: *args (arrays or trees of arrays): Each argument can be a single array or a tree of arrays. If a tree is given the nodes can be a Python :class:`list`, :class:`tuple` or :class:`dict`. Leaves which are not arrays are ignored. """ def exp(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise exponential. Args: a (array): Input array. Returns: array: The exponential of ``a``. """ def expand_dims( a: array, /, axis: int | Sequence[int], *, stream: Stream | Device | None = ..., ) -> array: """ Add a size one dimension at the given axis. Args: a (array): Input array. axes (int or tuple(int)): The index of the inserted dimensions. Returns: array: The array with inserted dimensions. """ def expm1(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise exponential minus 1. Computes ``exp(x) - 1`` with greater precision for small ``x``. Args: a (array): Input array. Returns: array: The expm1 of ``a``. """ def export_function( arg0: object, fun: Callable, *args, shapeless: bool = ..., **kwargs ) -> None: """ Export an MLX function. Example input arrays must be provided to export a function. The example inputs can be variable ``*args`` and ``**kwargs`` or a tuple of arrays and/or dictionary of string keys with array values. .. warning:: This is part of an experimental API which is likely to change in future versions of MLX. Functions exported with older versions of MLX may not be compatible with future versions. Args: file (str or Callable): Either a file path to export the function to or a callback. fun (Callable): A function which takes as input zero or more :class:`array` and returns one or more :class:`array`. *args (array): Example array inputs to the function. shapeless (bool, optional): Whether or not the function allows inputs with variable shapes. Default: ``False``. **kwargs (array): Additional example keyword array inputs to the function. Example: .. code-block:: python def fun(x, y): return x + y x = mx.array(1) y = mx.array([1, 2, 3]) mx.export_function("fun.mlxfn", fun, x, y=y) """ def export_to_dot(file: object, *args, **kwargs) -> None: """ Export a graph to DOT format for visualization. A variable number of output arrays can be provided for exporting The graph exported will recursively include all unevaluated inputs of the provided outputs. Args: file (str): The file path to export to. *args (array): The output arrays. **kwargs (dict[str, array]): Provide some names for arrays in the graph to make the result easier to parse. Example: >>> a = mx.array(1) + mx.array(2) >>> mx.export_to_dot("graph.dot", a) >>> x = mx.array(1) >>> y = mx.array(2) >>> mx.export_to_dot("graph.dot", x + y, x=x, y=y) """ def exporter(file: str, fun: Callable, *, shapeless: bool = ...) -> FunctionExporter: """ Make a callable object to export multiple traces of a function to a file. .. warning:: This is part of an experimental API which is likely to change in future versions of MLX. Functions exported with older versions of MLX may not be compatible with future versions. Args: file (str): File path to export the function to. shapeless (bool, optional): Whether or not the function allows inputs with variable shapes. Default: ``False``. Example: .. code-block:: python def fun(*args): return sum(args) with mx.exporter("fun.mlxfn", fun) as exporter: exporter(mx.array(1)) exporter(mx.array(1), mx.array(2)) exporter(mx.array(1), mx.array(2), mx.array(3)) """ def eye( n: int, m: int | None = ..., k: int = ..., dtype: Dtype | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Create an identity matrix or a general diagonal matrix. Args: n (int): The number of rows in the output. m (int, optional): The number of columns in the output. Defaults to n. k (int, optional): Index of the diagonal. Defaults to 0 (main diagonal). dtype (Dtype, optional): Data type of the output array. Defaults to float32. stream (Stream, optional): Stream or device. Defaults to None. Returns: array: An array where all elements are equal to zero, except for the k-th diagonal, whose values are equal to one. """ class finfo: """Get information on floating-point types.""" def __init__(self, arg: Dtype, /) -> None: ... @property def min(self) -> float: """The smallest representable number.""" @property def max(self) -> float: """The largest representable number.""" @property def eps(self) -> float: """ The difference between 1.0 and the next smallest representable number larger than 1.0. """ @property def dtype(self) -> Dtype: """The :obj:`Dtype`.""" def __repr__(self) -> str: ... def flatten( a: array, /, start_axis: int = ..., end_axis: int = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Flatten an array. The axes flattened will be between ``start_axis`` and ``end_axis``, inclusive. Negative axes are supported. After converting negative axis to positive, axes outside the valid range will be clamped to a valid value, ``start_axis`` to ``0`` and ``end_axis`` to ``ndim - 1``. Args: a (array): Input array. start_axis (int, optional): The first dimension to flatten. Defaults to ``0``. end_axis (int, optional): The last dimension to flatten. Defaults to ``-1``. stream (Stream, optional): Stream or device. Defaults to ``None`` in which case the default stream of the default device is used. Returns: array: The flattened array. Example: >>> a = mx.array([[1, 2], [3, 4]]) >>> mx.flatten(a) array([1, 2, 3, 4], dtype=int32) >>> >>> mx.flatten(a, start_axis=0, end_axis=-1) array([1, 2, 3, 4], dtype=int32) """ float16: Dtype = ... float32: Dtype = ... float64: Dtype = ... floating: DtypeCategory = ... def floor(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise floor. Args: a (array): Input array. Returns: array: The floor of ``a``. """ def floor_divide( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise integer division. If either array is a floating point type then it is equivalent to calling :func:`floor` after :func:`divide`. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The quotient ``a // b``. """ def full( shape: int | Sequence[int], vals: scalar | array, dtype: Dtype | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Construct an array with the given value. Constructs an array of size ``shape`` filled with ``vals``. If ``vals`` is an :obj:`array` it must be broadcastable to the given ``shape``. Args: shape (int or list(int)): The shape of the output array. vals (float or int or array): Values to fill the array with. dtype (Dtype, optional): Data type of the output array. If unspecified the output type is inferred from ``vals``. Returns: array: The output array with the specified shape and values. """ def gather_mm( a: array, b: array, /, lhs_indices: array, rhs_indices: array, *, sorted_indices: bool = ..., stream: Stream | Device | None = ..., ) -> array: """ Matrix multiplication with matrix-level gather. Performs a gather of the operands with the given indices followed by a (possibly batched) matrix multiplication of two arrays. This operation is more efficient than explicitly applying a :func:`take` followed by a :func:`matmul`. The indices ``lhs_indices`` and ``rhs_indices`` contain flat indices along the batch dimensions (i.e. all but the last two dimensions) of ``a`` and ``b`` respectively. For ``a`` with shape ``(A1, A2, ..., AS, M, K)``, ``lhs_indices`` contains indices from the range ``[0, A1 * A2 * ... * AS)`` For ``b`` with shape ``(B1, B2, ..., BS, M, K)``, ``rhs_indices`` contains indices from the range ``[0, B1 * B2 * ... * BS)`` If only one index is passed and it is sorted, the ``sorted_indices`` flag can be passed for a possible faster implementation. Args: a (array): Input array. b (array): Input array. lhs_indices (array, optional): Integer indices for ``a``. Default: ``None`` rhs_indices (array, optional): Integer indices for ``b``. Default: ``None`` sorted_indices (bool, optional): May allow a faster implementation if the passed indices are sorted. Default: ``False``. Returns: array: The output array. """ def gather_qmm( x: array, w: array, /, scales: array, biases: array | None = ..., lhs_indices: array | None = ..., rhs_indices: array | None = ..., transpose: bool = ..., group_size: int = ..., bits: int = ..., mode: str = ..., *, sorted_indices: bool = ..., stream: Stream | Device | None = ..., ) -> array: """ Perform quantized matrix multiplication with matrix-level gather. This operation is the quantized equivalent to :func:`gather_mm`. Similar to :func:`gather_mm`, the indices ``lhs_indices`` and ``rhs_indices`` contain flat indices along the batch dimensions (i.e. all but the last two dimensions) of ``x`` and ``w`` respectively. Note that ``scales`` and ``biases`` must have the same batch dimensions as ``w`` since they represent the same quantized matrix. Args: x (array): Input array w (array): Quantized matrix packed in unsigned integers scales (array): The scales to use per ``group_size`` elements of ``w`` biases (array, optional): The biases to use per ``group_size`` elements of ``w``. Default: ``None``. lhs_indices (array, optional): Integer indices for ``x``. Default: ``None``. rhs_indices (array, optional): Integer indices for ``w``. Default: ``None``. transpose (bool, optional): Defines whether to multiply with the transposed ``w`` or not, namely whether we are performing ``x @ w.T`` or ``x @ w``. Default: ``True``. group_size (int, optional): The size of the group in ``w`` that shares a scale and bias. Default: ``64``. bits (int, optional): The number of bits occupied by each element in ``w``. Default: ``4``. mode (str, optional): The quantization mode. Default: ``"affine"``. sorted_indices (bool, optional): May allow a faster implementation if the passed indices are sorted. Default: ``False``. Returns: array: The result of the multiplication of ``x`` with ``w`` after gathering using ``lhs_indices`` and ``rhs_indices``. """ generic: DtypeCategory = ... def get_active_memory() -> int: """ Get the actively used memory in bytes. Note, this will not always match memory use reported by the system because it does not include cached memory buffers. """ def get_cache_memory() -> int: """ Get the cache size in bytes. The cache includes memory not currently used that has not been returned to the system allocator. """ def get_peak_memory() -> int: """ Get the peak amount of used memory in bytes. The maximum memory used recorded from the beginning of the program execution or since the last call to :func:`reset_peak_memory`. """ gpu: DeviceType = ... def grad( fun: Callable, argnums: int | Sequence[int] | None = ..., argnames: str | Sequence[str] = ..., ) -> Callable: """ Returns a function which computes the gradient of ``fun``. Args: fun (Callable): A function which takes a variable number of :class:`array` or trees of :class:`array` and returns a scalar output :class:`array`. argnums (int or list(int), optional): Specify the index (or indices) of the positional arguments of ``fun`` to compute the gradient with respect to. If neither ``argnums`` nor ``argnames`` are provided ``argnums`` defaults to ``0`` indicating ``fun``'s first argument. argnames (str or list(str), optional): Specify keyword arguments of ``fun`` to compute gradients with respect to. It defaults to [] so no gradients for keyword arguments by default. Returns: Callable: A function which has the same input arguments as ``fun`` and returns the gradient(s). """ def greater( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise greater than. Strict greater than on two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The element-wise comparison ``a > b``. """ def greater_equal( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise greater or equal. Greater than or equal on two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The element-wise comparison ``a >= b``. """ def hadamard_transform( a: array, scale: float | None = ..., stream: Stream | Device | None = ... ) -> array: """ Perform the Walsh-Hadamard transform along the final axis. Equivalent to: .. code-block:: python from scipy.linalg import hadamard y = (hadamard(len(x)) @ x) * scale Supports sizes ``n = m*2^k`` for ``m`` in ``(1, 12, 20, 28)`` and ``2^k <= 8192`` for float32 and ``2^k <= 16384`` for float16/bfloat16. Args: a (array): Input array or scalar. scale (float): Scale the output by this factor. Defaults to ``1/sqrt(a.shape[-1])`` so that the Hadamard matrix is orthonormal. Returns: array: The transformed array. """ def identity( n: int, dtype: Dtype | None = ..., *, stream: Stream | Device | None = ... ) -> array: """ Create a square identity matrix. Args: n (int): The number of rows and columns in the output. dtype (Dtype, optional): Data type of the output array. Defaults to float32. stream (Stream, optional): Stream or device. Defaults to None. Returns: array: An identity matrix of size n x n. """ class iinfo: """Get information on integer types.""" def __init__(self, arg: Dtype, /) -> None: ... @property def min(self) -> int: """The smallest representable number.""" @property def max(self) -> int: """The largest representable number.""" @property def dtype(self) -> Dtype: """The :obj:`Dtype`.""" def __repr__(self) -> str: ... def imag(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Returns the imaginary part of a complex array. Args: a (array): Input array. Returns: array: The imaginary part of ``a``. """ def import_function(file: str) -> Callable: """ Import a function from a file. The imported function can be called either with ``*args`` and ``**kwargs`` or with a tuple of arrays and/or dictionary of string keys with array values. Imported functions always return a tuple of arrays. .. warning:: This is part of an experimental API which is likely to change in future versions of MLX. Functions exported with older versions of MLX may not be compatible with future versions. Args: file (str): The file path to import the function from. Returns: Callable: The imported function. Example: >>> fn = mx.import_function("function.mlxfn") >>> out = fn(a, b, x=x, y=y)[0] >>> >>> out = fn((a, b), {"x": x, "y": y}[0] """ inexact: DtypeCategory = ... inf: float = ... def inner(a: array, b: array, /, *, stream: Stream | Device | None = ...) -> array: """ Ordinary inner product of vectors for 1-D arrays, in higher dimensions a sum product over the last axes. Args: a (array): Input array b (array): Input array Returns: array: The inner product. """ int16: Dtype = ... int32: Dtype = ... int64: Dtype = ... int8: Dtype = ... integer: DtypeCategory = ... def is_available(device: Device) -> bool: """Check if a back-end is available for the given device.""" def isclose( a: array, b: array, /, rtol: float = ..., atol: float = ..., *, equal_nan: bool = ..., stream: Stream | Device | None = ..., ) -> array: """ Returns a boolean array where two arrays are element-wise equal within a tolerance. Infinite values are considered equal if they have the same sign, NaN values are not equal unless ``equal_nan`` is ``True``. Two values are considered equal if: .. code-block:: abs(a - b) <= (atol + rtol * abs(b)) Note unlike :func:`array_equal`, this function supports numpy-style broadcasting. Args: a (array): Input array. b (array): Input array. rtol (float): Relative tolerance. atol (float): Absolute tolerance. equal_nan (bool): If ``True``, NaNs are considered equal. Defaults to ``False``. Returns: array: The boolean output scalar indicating if the arrays are close. """ def isfinite(a: array, stream: Stream | Device | None = ...) -> array: """ Return a boolean array indicating which elements are finite. An element is finite if it is not infinite or NaN. Args: a (array): Input array. Returns: array: The boolean array indicating which elements are finite. """ def isinf(a: array, stream: Stream | Device | None = ...) -> array: """ Return a boolean array indicating which elements are +/- inifnity. Args: a (array): Input array. Returns: array: The boolean array indicating which elements are +/- infinity. """ def isnan(a: array, stream: Stream | Device | None = ...) -> array: """ Return a boolean array indicating which elements are NaN. Args: a (array): Input array. Returns: array: The boolean array indicating which elements are NaN. """ def isneginf(a: array, stream: Stream | Device | None = ...) -> array: """ Return a boolean array indicating which elements are negative infinity. Args: a (array): Input array. stream (Stream | Device | None): Optional stream or device. Returns: array: The boolean array indicating which elements are negative infinity. """ def isposinf(a: array, stream: Stream | Device | None = ...) -> array: """ Return a boolean array indicating which elements are positive infinity. Args: a (array): Input array. stream (Stream | Device | None): Optional stream or device. Returns: array: The boolean array indicating which elements are positive infinity. """ def issubdtype(arg1: Dtype | DtypeCategory, arg2: Dtype | DtypeCategory) -> bool: """ Check if a :obj:`Dtype` or :obj:`DtypeCategory` is a subtype of another. Args: arg1 (Dtype | DtypeCategory: First dtype or category. arg2 (Dtype | DtypeCategory: Second dtype or category. Returns: bool: A boolean indicating if the first input is a subtype of the second input. Example: >>> ints = mx.array([1, 2, 3], dtype=mx.int32) >>> mx.issubdtype(ints.dtype, mx.integer) True >>> mx.issubdtype(ints.dtype, mx.floating) False >>> floats = mx.array([1, 2, 3], dtype=mx.float32) >>> mx.issubdtype(floats.dtype, mx.integer) False >>> mx.issubdtype(floats.dtype, mx.floating) True Similar types of different sizes are not subdtypes of each other: >>> mx.issubdtype(mx.float64, mx.float32) False >>> mx.issubdtype(mx.float32, mx.float64) False but both are subtypes of `floating`: >>> mx.issubdtype(mx.float64, mx.floating) True >>> mx.issubdtype(mx.float32, mx.floating) True For convenience, dtype-like objects are allowed too: >>> mx.issubdtype(mx.float32, mx.inexact) True >>> mx.issubdtype(mx.signedinteger, mx.floating) False """ def jvp( fun: Callable, primals: list[array], tangents: list[array] ) -> tuple[list[array], list[array]]: """ Compute the Jacobian-vector product. This computes the product of the Jacobian of a function ``fun`` evaluated at ``primals`` with the ``tangents``. Args: fun (Callable): A function which takes a variable number of :class:`array` and returns a single :class:`array` or list of :class:`array`. primals (list(array)): A list of :class:`array` at which to evaluate the Jacobian. tangents (list(array)): A list of :class:`array` which are the "vector" in the Jacobian-vector product. The ``tangents`` should be the same in number, shape, and type as the inputs of ``fun`` (i.e. the ``primals``). Returns: list(array): A list of the Jacobian-vector products which is the same in number, shape, and type of the inputs to ``fun``. """ def kron(a: array, b: array, *, stream: Stream | Device | None = ...) -> array: """ Compute the Kronecker product of two arrays ``a`` and ``b``. Args: a (array): The first input array. b (array): The second input array. stream (Stream | Device | None, optional): Optional stream or device for execution. Default: ``None``. Returns: array: The Kronecker product of ``a`` and ``b``. Examples: >>> a = mx.array([[1, 2], [3, 4]]) >>> b = mx.array([[0, 5], [6, 7]]) >>> result = mx.kron(a, b) >>> print(result) array([[0, 5, 0, 10], [6, 7, 12, 14], [0, 15, 0, 20], [18, 21, 24, 28]], dtype=int32) """ def left_shift( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise left shift. Shift the bits of the first input to the left by the second using numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The bitwise left shift ``a << b``. """ def less( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise less than. Strict less than on two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The element-wise comparison ``a < b``. """ def less_equal( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise less than or equal. Less than or equal on two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The element-wise comparison ``a <= b``. """ def linspace( start, stop, num: int | None = ..., dtype: Dtype | None = ..., stream: Stream | Device | None = ..., ) -> array: """ Generate ``num`` evenly spaced numbers over interval ``[start, stop]``. Args: start (scalar): Starting value. stop (scalar): Stopping value. num (int, optional): Number of samples, defaults to ``50``. dtype (Dtype, optional): Specifies the data type of the output, default to ``float32``. Returns: array: The range of values. """ def load( file: str | pathlib.Path, /, format: str | None = ..., return_metadata: bool = ..., *, stream: Stream | Device | None = ..., ) -> array | dict[str, array]: """ Load array(s) from a binary file. The supported formats are ``.npy``, ``.npz``, ``.safetensors``, and ``.gguf``. Args: file (str, pathlib.Path): File in which the array is saved. format (str, optional): Format of the file. If ``None``, the format is inferred from the file extension. Supported formats: ``npy``, ``npz``, and ``safetensors``. Default: ``None``. return_metadata (bool, optional): Load the metadata for formats which support matadata. The metadata will be returned as an additional dictionary. Default: ``False``. Returns: array or dict: A single array if loading from a ``.npy`` file or a dict mapping names to arrays if loading from a ``.npz`` or ``.safetensors`` file. If ``return_metadata`` is ``True`` an additional dictionary of metadata will be returned. Warning: When loading unsupported quantization formats from GGUF, tensors will automatically cast to ``mx.float16`` """ def log(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise natural logarithm. Args: a (array): Input array. Returns: array: The natural logarithm of ``a``. """ def log10(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise base-10 logarithm. Args: a (array): Input array. Returns: array: The base-10 logarithm of ``a``. """ def log1p(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise natural log of one plus the array. Args: a (array): Input array. Returns: array: The natural logarithm of one plus ``a``. """ def log2(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise base-2 logarithm. Args: a (array): Input array. Returns: array: The base-2 logarithm of ``a``. """ def logaddexp( a: scalar | array, b: scalar | array, /, *, stream: Stream | Device | None = ..., ) -> array: """ Element-wise log-add-exp. This is a numerically stable log-add-exp of two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. The computation is is a numerically stable version of ``log(exp(a) + exp(b))``. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The log-add-exp of ``a`` and ``b``. """ def logcumsumexp( a: array, /, axis: int | None = ..., *, reverse: bool = ..., inclusive: bool = ..., stream: Stream | Device | None = ..., ) -> array: """ Return the cumulative logsumexp of the elements along the given axis. Args: a (array): Input array axis (int, optional): Optional axis to compute the cumulative logsumexp over. If unspecified the cumulative logsumexp of the flattened array is returned. reverse (bool): Perform the cumulative logsumexp in reverse. inclusive (bool): The i-th element of the output includes the i-th element of the input. Returns: array: The output array. """ def logical_and( a: array, b: array, /, *, stream: Stream | Device | None = ... ) -> array: """ Element-wise logical and. Args: a (array): First input array or scalar. b (array): Second input array or scalar. Returns: array: The boolean array containing the logical and of ``a`` and ``b``. """ def logical_not(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise logical not. Args: a (array): Input array or scalar. Returns: array: The boolean array containing the logical not of ``a``. """ def logical_or(a: array, b: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise logical or. Args: a (array): First input array or scalar. b (array): Second input array or scalar. Returns: array: The boolean array containing the logical or of ``a`` and ``b``. """ def logsumexp( a: array, /, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """ A `log-sum-exp` reduction over the given axes. The log-sum-exp reduction is a numerically stable version of: .. code-block:: log(sum(exp(a), axis)) Args: a (array): Input array. axis (int or list(int), optional): Optional axis or axes to reduce over. If unspecified this defaults to reducing over the entire array. keepdims (bool, optional): Keep reduced axes as singleton dimensions, defaults to `False`. Returns: array: The output array with the corresponding axes reduced. """ def matmul(a: array, b: array, /, *, stream: Stream | Device | None = ...) -> array: """ Matrix multiplication. Perform the (possibly batched) matrix multiplication of two arrays. This function supports broadcasting for arrays with more than two dimensions. - If the first array is 1-D then a 1 is prepended to its shape to make it a matrix. Similarly if the second array is 1-D then a 1 is appended to its shape to make it a matrix. In either case the singleton dimension is removed from the result. - A batched matrix multiplication is performed if the arrays have more than 2 dimensions. The matrix dimensions for the matrix product are the last two dimensions of each input. - All but the last two dimensions of each input are broadcast with one another using standard numpy-style broadcasting semantics. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The matrix product of ``a`` and ``b``. """ def max( a: array, /, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """ A `max` reduction over the given axes. Args: a (array): Input array. axis (int or list(int), optional): Optional axis or axes to reduce over. If unspecified this defaults to reducing over the entire array. keepdims (bool, optional): Keep reduced axes as singleton dimensions, defaults to `False`. Returns: array: The output array with the corresponding axes reduced. """ def maximum( a: scalar | array, b: scalar | array, /, *, stream: Stream | Device | None = ..., ) -> array: """ Element-wise maximum. Take the element-wise max of two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The max of ``a`` and ``b``. """ def mean( a: array, /, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Compute the mean(s) over the given axes. Args: a (array): Input array. axis (int or list(int), optional): Optional axis or axes to reduce over. If unspecified this defaults to reducing over the entire array. keepdims (bool, optional): Keep reduced axes as singleton dimensions, defaults to `False`. Returns: array: The output array of means. """ def meshgrid( *arrays: array, sparse: bool | None = ..., indexing: str | None = ..., stream: Stream | Device | None = ..., ) -> array: """ Generate multidimensional coordinate grids from 1-D coordinate arrays Args: *arrays (array): Input arrays. sparse (bool, optional): If ``True``, a sparse grid is returned in which each output array has a single non-zero element. If ``False``, a dense grid is returned. Defaults to ``False``. indexing (str, optional): Cartesian ('xy') or matrix ('ij') indexing of the output arrays. Defaults to ``'xy'``. Returns: list(array): The output arrays. """ def min( a: array, /, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """ A `min` reduction over the given axes. Args: a (array): Input array. axis (int or list(int), optional): Optional axis or axes to reduce over. If unspecified this defaults to reducing over the entire array. keepdims (bool, optional): Keep reduced axes as singleton dimensions, defaults to `False`. Returns: array: The output array with the corresponding axes reduced. """ def minimum( a: scalar | array, b: scalar | array, /, *, stream: Stream | Device | None = ..., ) -> array: """ Element-wise minimum. Take the element-wise min of two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The min of ``a`` and ``b``. """ def moveaxis( a: array, /, source: int, destination: int, *, stream: Stream | Device | None = ..., ) -> array: """ Move an axis to a new position. Args: a (array): Input array. source (int): Specifies the source axis. destination (int): Specifies the destination axis. Returns: array: The array with the axis moved. """ def multiply( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise multiplication. Multiply two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The multiplication ``a * b``. """ nan: float = ... def nan_to_num( a: scalar | array, nan: float = ..., posinf: float | None = ..., neginf: float | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Replace NaN and Inf values with finite numbers. Args: a (array): Input array nan (float, optional): Value to replace NaN with. Default: ``0``. posinf (float, optional): Value to replace positive infinities with. If ``None``, defaults to largest finite value for the given data type. Default: ``None``. neginf (float, optional): Value to replace negative infinities with. If ``None``, defaults to the negative of the largest finite value for the given data type. Default: ``None``. Returns: array: Output array with NaN and Inf replaced. """ def negative(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise negation. Args: a (array): Input array. Returns: array: The negative of ``a``. """ def new_stream(device: Device) -> Stream: """Make a new stream on the given device.""" newaxis: None = ... def not_equal( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise not equal. Not equal comparison on two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The element-wise comparison ``a != b``. """ number: DtypeCategory = ... def ones( shape: int | Sequence[int], dtype: Dtype | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Construct an array of ones. Args: shape (int or list(int)): The shape of the output array. dtype (Dtype, optional): Data type of the output array. If unspecified the output type defaults to ``float32``. Returns: array: The array of ones with the specified shape. """ def ones_like(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ An array of ones like the input. Args: a (array): The input to take the shape and type from. Returns: array: The output array filled with ones. """ def outer(a: array, b: array, /, *, stream: Stream | Device | None = ...) -> array: """ Compute the outer product of two 1-D arrays, if the array's passed are not 1-D a flatten op will be run beforehand. Args: a (array): Input array b (array): Input array Returns: array: The outer product. """ def pad( a: array, pad_width: int | tuple[int] | tuple[int, int] | list[tuple[int, int]], mode: Literal["constant", "edge"] = ..., constant_values: scalar | array = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Pad an array with a constant value Args: a (array): Input array. pad_width (int, tuple(int), tuple(int, int) or list(tuple(int, int))): Number of padded values to add to the edges of each axis:``((before_1, after_1), (before_2, after_2), ..., (before_N, after_N))``. If a single pair of integers is passed then ``(before_i, after_i)`` are all the same. If a single integer or tuple with a single integer is passed then all axes are extended by the same number on each side. mode: Padding mode. One of the following strings: "constant" (default): Pads with a constant value. "edge": Pads with the edge values of array. constant_value (array or scalar, optional): Optional constant value to pad the edges of the array with. Returns: array: The padded array. """ def partition( a: array, /, kth: int, axis: int | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Returns a partitioned copy of the array such that the smaller ``kth`` elements are first. The ordering of the elements in partitions is undefined. Args: a (array): Input array. kth (int): Element at the ``kth`` index will be in its sorted position in the output. All elements before the kth index will be less or equal to the ``kth`` element and all elements after will be greater or equal to the ``kth`` element in the output. axis (int or None, optional): Optional axis to partition over. If ``None``, this partitions over the flattened array. If unspecified, it defaults to ``-1``. Returns: array: The partitioned array. """ def permute_dims( a: array, /, axes: Sequence[int] | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """See :func:`transpose`.""" pi: float = ... def power( a: scalar | array, b: scalar | array, /, *, stream: Stream | Device | None = ..., ) -> array: """ Element-wise power operation. Raise the elements of a to the powers in elements of b with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: Bases of ``a`` raised to powers in ``b``. """ def prod( a: array, /, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """ An product reduction over the given axes. Args: a (array): Input array. axis (int or list(int), optional): Optional axis or axes to reduce over. If unspecified this defaults to reducing over the entire array. keepdims (bool, optional): Keep reduced axes as singleton dimensions, defaults to `False`. Returns: array: The output array with the corresponding axes reduced. """ def put_along_axis( a: array, /, indices: array, values: array, axis: int | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Put values along an axis at the specified indices. Args: a (array): Destination array. indices (array): Indices array. These should be broadcastable with the input array excluding the `axis` dimension. values (array): Values array. These should be broadcastable with the indices. axis (int or None): Axis in the destination to put the values to. If ``axis == None`` the destination is flattened prior to the put operation. Returns: array: The output array. """ def quantize( w: array, /, group_size: int = ..., bits: int = ..., mode: str = ..., *, stream: Stream | Device | None = ..., ) -> tuple[array, array, array]: r""" Quantize the matrix ``w`` using ``bits`` bits per element. Note, every ``group_size`` elements in a row of ``w`` are quantized together. Hence, number of columns of ``w`` should be divisible by ``group_size``. In particular, the rows of ``w`` are divided into groups of size ``group_size`` which are quantized together. .. warning:: ``quantize`` currently only supports 2D inputs with the second dimension divisible by ``group_size`` The supported quantization modes are ``"affine"`` and ``"mxfp4"``. They are described in more detail below. Args: w (array): Matrix to be quantized group_size (int, optional): The size of the group in ``w`` that shares a scale and bias. Default: ``64``. bits (int, optional): The number of bits occupied by each element of ``w`` in the returned quantized matrix. Default: ``4``. mode (str, optional): The quantization mode. Default: ``"affine"``. Returns: tuple: A tuple with either two or three elements containing: * w_q (array): The quantized version of ``w`` * scales (array): The quantization scales * biases (array): The quantization biases (returned for ``mode=="affine"``). Notes: The ``affine`` mode quantizes groups of :math:`g` consecutive elements in a row of ``w``. For each group the quantized representation of each element :math:`\hat{w_i}` is computed as follows: .. math:: \begin{aligned} \alpha &= \max_i w_i \\ \beta &= \min_i w_i \\ s &= \frac{\alpha - \beta}{2^b - 1} \\ \hat{w_i} &= \textrm{round}\left( \frac{w_i - \beta}{s}\right). \end{aligned} After the above computation, :math:`\hat{w_i}` fits in :math:`b` bits and is packed in an unsigned 32-bit integer from the lower to upper bits. For instance, for 4-bit quantization we fit 8 elements in an unsigned 32 bit integer where the 1st element occupies the 4 least significant bits, the 2nd bits 4-7 etc. To dequantize the elements of ``w``, we also save :math:`s` and :math:`\beta` which are the returned ``scales`` and ``biases`` respectively. The ``mxfp4`` mode similarly quantizes groups of :math:`g` elements of ``w``. For ``mxfp4`` the group size must be ``32``. The elements are quantized to 4-bit precision floating-point values (E2M1) with a shared 8-bit scale per group. Unlike ``affine`` quantization, ``mxfp4`` does not have a bias value. More details on the format can be found in the `specification `_. """ def quantized_matmul( x: array, w: array, /, scales: array, biases: array | None = ..., transpose: bool = ..., group_size: int = ..., bits: int = ..., mode: str = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Perform the matrix multiplication with the quantized matrix ``w``. The quantization uses one floating point scale and bias per ``group_size`` of elements. Each element in ``w`` takes ``bits`` bits and is packed in an unsigned 32 bit integer. Args: x (array): Input array w (array): Quantized matrix packed in unsigned integers scales (array): The scales to use per ``group_size`` elements of ``w`` biases (array, optional): The biases to use per ``group_size`` elements of ``w``. Default: ``None``. transpose (bool, optional): Defines whether to multiply with the transposed ``w`` or not, namely whether we are performing ``x @ w.T`` or ``x @ w``. Default: ``True``. group_size (int, optional): The size of the group in ``w`` that shares a scale and bias. Default: ``64``. bits (int, optional): The number of bits occupied by each element in ``w``. Default: ``4``. mode (str, optional): The quantization mode. Default: ``"affine"``. Returns: array: The result of the multiplication of ``x`` with ``w``. """ def radians(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Convert angles from degrees to radians. Args: a (array): Input array. Returns: array: The angles in radians. """ def real(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Returns the real part of a complex array. Args: a (array): Input array. Returns: array: The real part of ``a``. """ def reciprocal(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise reciprocal. Args: a (array): Input array. Returns: array: The reciprocal of ``a``. """ def remainder( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise remainder of division. Computes the remainder of dividing a with b with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The remainder of ``a // b``. """ def repeat( array: array, repeats: int, axis: int | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Repeat an array along a specified axis. Args: array (array): Input array. repeats (int): The number of repetitions for each element. axis (int, optional): The axis in which to repeat the array along. If unspecified it uses the flattened array of the input and repeats along axis 0. stream (Stream, optional): Stream or device. Defaults to ``None``. Returns: array: The resulting repeated array. """ def reset_peak_memory() -> None: """Reset the peak memory to zero.""" def reshape( a: array, /, shape: Sequence[int], *, stream: Stream | Device | None = ... ) -> array: """ Reshape an array while preserving the size. Args: a (array): Input array. shape (tuple(int)): New shape. stream (Stream, optional): Stream or device. Defaults to ``None`` in which case the default stream of the default device is used. Returns: array: The reshaped array. """ def right_shift( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise right shift. Shift the bits of the first input to the right by the second using numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The bitwise right shift ``a >> b``. """ def roll( a: array, shift: int | tuple[int], axis: int | tuple[int] | None = ..., /, *, stream: Stream | Device | None = ..., ) -> array: """ Roll array elements along a given axis. Elements that are rolled beyond the end of the array are introduced at the beggining and vice-versa. If the axis is not provided the array is flattened, rolled and then the shape is restored. Args: a (array): Input array shift (int or tuple(int)): The number of places by which elements are shifted. If positive the array is rolled to the right, if negative it is rolled to the left. If an int is provided but the axis is a tuple then the same value is used for all axes. axis (int or tuple(int), optional): The axis or axes along which to roll the elements. """ def round( a: array, /, decimals: int = ..., stream: Stream | Device | None = ... ) -> array: """ Round to the given number of decimals. Basically performs: .. code-block:: python s = 10**decimals x = round(x * s) / s Args: a (array): Input array decimals (int): Number of decimal places to round to. (default: 0) Returns: array: An array of the same type as ``a`` rounded to the given number of decimals. """ def rsqrt(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise reciprocal and square root. Args: a (array): Input array. Returns: array: One over the square root of ``a``. """ def save(file: str | pathlib.Path, arr: array) -> None: """ Save the array to a binary file in ``.npy`` format. Args: file (str, pathlib.Path): File to which the array is saved arr (array): Array to be saved. """ def save_gguf( file: str | pathlib.Path, arrays: dict[str, array], metadata: dict[str, array | str | list[str]], ): """ Save array(s) to a binary file in ``.gguf`` format. See the `GGUF documentation `_ for more information on the format. Args: file (file, str, pathlib.Path): File in which the array is saved. arrays (dict(str, array)): The dictionary of names to arrays to be saved. metadata (dict(str, array | str | list(str))): The dictionary of metadata to be saved. The values can be a scalar or 1D obj:`array`, a :obj:`str`, or a :obj:`list` of :obj:`str`. """ def save_safetensors( file: str | pathlib.Path, arrays: dict[str, array], metadata: dict[str, str] | None = ..., ): """ Save array(s) to a binary file in ``.safetensors`` format. See the `Safetensors documentation `_ for more information on the format. Args: file (file, str, pathlib.Path): File in which the array is saved. arrays (dict(str, array)): The dictionary of names to arrays to be saved. metadata (dict(str, str), optional): The dictionary of metadata to be saved. """ def savez(file: str | pathlib.Path, *args, **kwargs): """ Save several arrays to a binary file in uncompressed ``.npz`` format. .. code-block:: python import mlx.core as mx x = mx.ones((10, 10)) mx.savez("my_path.npz", x=x) import mlx.nn as nn from mlx.utils import tree_flatten model = nn.TransformerEncoder(6, 128, 4) flat_params = tree_flatten(model.parameters()) mx.savez("model.npz", **dict(flat_params)) Args: file (file, str, pathlib.Path): Path to file to which the arrays are saved. *args (arrays): Arrays to be saved. **kwargs (arrays): Arrays to be saved. Each array will be saved with the associated keyword as the output file name. """ def savez_compressed(file: str | pathlib.Path, *args, **kwargs): """ Save several arrays to a binary file in compressed ``.npz`` format. Args: file (file, str, pathlib.Path): Path to file to which the arrays are saved. *args (arrays): Arrays to be saved. **kwargs (arrays): Arrays to be saved. Each array will be saved with the associated keyword as the output file name. """ def segmented_mm( a: array, b: array, /, segments: array, *, stream: Stream | Device | None = ... ) -> array: """ Perform a matrix multiplication but segment the inner dimension and save the result for each segment separately. Args: a (array): Input array of shape ``MxK``. b (array): Input array of shape ``KxN``. segments (array): The offsets into the inner dimension for each segment. Returns: array: The result per segment of shape ``MxN``. """ def set_cache_limit(limit: int) -> int: """ Set the free cache limit. If using more than the given limit, free memory will be reclaimed from the cache on the next allocation. To disable the cache, set the limit to ``0``. The cache limit defaults to the memory limit. See :func:`set_memory_limit` for more details. Args: limit (int): The cache limit in bytes. Returns: int: The previous cache limit in bytes. """ def set_default_device(device: Device | DeviceType) -> None: """Set the default device.""" def set_default_stream(stream: Stream) -> None: """ Set the default stream. This will make the given stream the default for the streams device. It will not change the default device. Args: stream (stream): Stream to make the default. """ def set_memory_limit(limit: int) -> int: """ Set the memory limit. The memory limit is a guideline for the maximum amount of memory to use during graph evaluation. If the memory limit is exceeded and there is no more RAM (including swap when available) allocations will result in an exception. When metal is available the memory limit defaults to 1.5 times the maximum recommended working set size reported by the device. Args: limit (int): Memory limit in bytes. Returns: int: The previous memory limit in bytes. """ def set_wired_limit(limit: int) -> int: """ Set the wired size limit. .. note:: * This function is only useful on macOS 15.0 or higher. * The wired limit should remain strictly less than the total memory size. The wired limit is the total size in bytes of memory that will be kept resident. The default value is ``0``. Setting a wired limit larger than system wired limit is an error. You can increase the system wired limit with: .. code-block:: sudo sysctl iogpu.wired_limit_mb= Use :func:`device_info` to query the system wired limit (``"max_recommended_working_set_size"``) and the total memory size (``"memory_size"``). Args: limit (int): The wired limit in bytes. Returns: int: The previous wired limit in bytes. """ def sigmoid(a: array, /, *, stream: Stream | Device | None = ...) -> array: r""" Element-wise logistic sigmoid. The logistic sigmoid function is: .. math:: \mathrm{sigmoid}(x) = \frac{1}{1 + e^{-x}} Args: a (array): Input array. Returns: array: The logistic sigmoid of ``a``. """ def sign(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise sign. Args: a (array): Input array. Returns: array: The sign of ``a``. """ signedinteger: DtypeCategory = ... def sin(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise sine. Args: a (array): Input array. Returns: array: The sine of ``a``. """ def sinh(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise hyperbolic sine. Args: a (array): Input array. Returns: array: The hyperbolic sine of ``a``. """ def slice( a: array, start_indices: array, axes: Sequence[int], slice_size: Sequence[int], *, stream: Stream | Device | None = ..., ) -> array: """ Extract a sub-array from the input array. Args: a (array): Input array start_indices (array): The index location to start the slice at. axes (tuple(int)): The axes corresponding to the indices in ``start_indices``. slice_size (tuple(int)): The size of the slice. Returns: array: The sliced output array. Example: >>> a = mx.array([[1, 2, 3], [4, 5, 6]]) >>> mx.slice(a, start_indices=mx.array(1), axes=(0,), slice_size=(1, 2)) array([[4, 5]], dtype=int32) >>> >>> mx.slice(a, start_indices=mx.array(1), axes=(1,), slice_size=(2, 1)) array([[2], [5]], dtype=int32) """ def slice_update( a: array, update: array, start_indices: array, axes: Sequence[int], *, stream: Stream | Device | None = ..., ) -> array: """ Update a sub-array of the input array. Args: a (array): The input array to update update (array): The update array. start_indices (array): The index location to start the slice at. axes (tuple(int)): The axes corresponding to the indices in ``start_indices``. Returns: array: The output array with the same shape and type as the input. Example: >>> a = mx.zeros((3, 3)) >>> mx.slice_update(a, mx.ones((1, 2)), start_indices=mx.array(1, 1), axes=(0, 1)) array([[0, 0, 0], [0, 1, 0], [0, 1, 0]], dtype=float32) """ def softmax( a: array, /, axis: int | Sequence[int] | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Perform the softmax along the given axis. This operation is a numerically stable version of: .. code-block:: exp(a) / sum(exp(a), axis, keepdims=True) Args: a (array): Input array. axis (int or list(int), optional): Optional axis or axes to compute the softmax over. If unspecified this performs the softmax over the full array. Returns: array: The output of the softmax. """ def sort( a: array, /, axis: int | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Returns a sorted copy of the array. Args: a (array): Input array. axis (int or None, optional): Optional axis to sort over. If ``None``, this sorts over the flattened array. If unspecified, it defaults to -1 (sorting over the last axis). Returns: array: The sorted array. """ def split( a: array, /, indices_or_sections: int | Sequence[int], axis: int = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Split an array along a given axis. Args: a (array): Input array. indices_or_sections (int or list(int)): If ``indices_or_sections`` is an integer the array is split into that many sections of equal size. An error is raised if this is not possible. If ``indices_or_sections`` is a list, the list contains the indices of the start of each subarray along the given axis. axis (int, optional): Axis to split along, defaults to `0`. Returns: list(array): A list of split arrays. """ def sqrt(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise square root. Args: a (array): Input array. Returns: array: The square root of ``a``. """ def square(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise square. Args: a (array): Input array. Returns: array: The square of ``a``. """ def squeeze( a: array, /, axis: int | Sequence[int] | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Remove length one axes from an array. Args: a (array): Input array. axis (int or tuple(int), optional): Axes to remove. Defaults to ``None`` in which case all size one axes are removed. Returns: array: The output array with size one axes removed. """ def stack( arrays: list[array], axis: int | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Stacks the arrays along a new axis. Args: arrays (list(array)): A list of arrays to stack. axis (int, optional): The axis in the result array along which the input arrays are stacked. Defaults to ``0``. stream (Stream, optional): Stream or device. Defaults to ``None``. Returns: array: The resulting stacked array. """ def std( a: array, /, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., ddof: int = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Compute the standard deviation(s) over the given axes. Args: a (array): Input array. axis (int or list(int), optional): Optional axis or axes to reduce over. If unspecified this defaults to reducing over the entire array. keepdims (bool, optional): Keep reduced axes as singleton dimensions, defaults to `False`. ddof (int, optional): The divisor to compute the variance is ``N - ddof``, defaults to 0. Returns: array: The output array of standard deviations. """ def stop_gradient(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Stop gradients from being computed. The operation is the identity but it prevents gradients from flowing through the array. Args: a (array): Input array. Returns: array: The unchanged input ``a`` but without gradient flowing through it. """ def stream(s: Stream | Device) -> StreamContext: """ Create a context manager to set the default device and stream. Args: s: The :obj:`Stream` or :obj:`Device` to set as the default. Returns: A context manager that sets the default device and stream. Example: .. code-block::python import mlx.core as mx # Create a context manager for the default device and stream. with mx.stream(mx.cpu): # Operations here will use mx.cpu by default. pass """ def subtract( a: scalar | array, b: scalar | array, stream: Stream | Device | None = ..., ) -> array: """ Element-wise subtraction. Subtract one array from another with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. Args: a (array): Input array or scalar. b (array): Input array or scalar. Returns: array: The difference ``a - b``. """ def sum( a: array, /, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Sum reduce the array over the given axes. Args: a (array): Input array. axis (int or list(int), optional): Optional axis or axes to reduce over. If unspecified this defaults to reducing over the entire array. keepdims (bool, optional): Keep reduced axes as singleton dimensions, defaults to `False`. Returns: array: The output array with the corresponding axes reduced. """ def swapaxes( a: array, /, axis1: int, axis2: int, *, stream: Stream | Device | None = ... ) -> array: """ Swap two axes of an array. Args: a (array): Input array. axis1 (int): Specifies the first axis. axis2 (int): Specifies the second axis. Returns: array: The array with swapped axes. """ def synchronize(stream: Stream | None = ...) -> None: """ Synchronize with the given stream. Args: stream (Stream, optional): The stream to synchronize with. If ``None`` then the default stream of the default device is used. Default: ``None``. """ def take( a: array, /, indices: int | array, axis: int | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Take elements along an axis. The elements are taken from ``indices`` along the specified axis. If the axis is not specified the array is treated as a flattened 1-D array prior to performing the take. As an example, if the ``axis=1`` this is equivalent to ``a[:, indices, ...]``. Args: a (array): Input array. indices (int or array): Integer index or input array with integral type. axis (int, optional): Axis along which to perform the take. If unspecified the array is treated as a flattened 1-D vector. Returns: array: The indexed values of ``a``. """ def take_along_axis( a: array, /, indices: array, axis: int | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Take values along an axis at the specified indices. Args: a (array): Input array. indices (array): Indices array. These should be broadcastable with the input array excluding the `axis` dimension. axis (int or None): Axis in the input to take the values from. If ``axis == None`` the array is flattened to 1D prior to the indexing operation. Returns: array: The output array. """ def tan(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise tangent. Args: a (array): Input array. Returns: array: The tangent of ``a``. """ def tanh(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ Element-wise hyperbolic tangent. Args: a (array): Input array. Returns: array: The hyperbolic tangent of ``a``. """ def tensordot( a: array, b: array, /, axes: int | list[Sequence[int]] = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Compute the tensor dot product along the specified axes. Args: a (array): Input array b (array): Input array axes (int or list(list(int)), optional): The number of dimensions to sum over. If an integer is provided, then sum over the last ``axes`` dimensions of ``a`` and the first ``axes`` dimensions of ``b``. If a list of lists is provided, then sum over the corresponding dimensions of ``a`` and ``b``. Default: 2. Returns: array: The tensor dot product. """ def tile( a: array, reps: int | Sequence[int], /, *, stream: Stream | Device | None = ..., ) -> array: """ Construct an array by repeating ``a`` the number of times given by ``reps``. Args: a (array): Input array reps (int or list(int)): The number of times to repeat ``a`` along each axis. Returns: array: The tiled array. """ def topk( a: array, /, k: int, axis: int | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Returns the ``k`` largest elements from the input along a given axis. The elements will not necessarily be in sorted order. Args: a (array): Input array. k (int): ``k`` top elements to be returned axis (int or None, optional): Optional axis to select over. If ``None``, this selects the top ``k`` elements over the flattened array. If unspecified, it defaults to ``-1``. Returns: array: The top ``k`` elements from the input. """ def trace( a: array, /, offset: int = ..., axis1: int = ..., axis2: int = ..., dtype: Dtype | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Return the sum along a specified diagonal in the given array. Args: a (array): Input array offset (int, optional): Offset of the diagonal from the main diagonal. Can be positive or negative. Default: ``0``. axis1 (int, optional): The first axis of the 2-D sub-arrays from which the diagonals should be taken. Default: ``0``. axis2 (int, optional): The second axis of the 2-D sub-arrays from which the diagonals should be taken. Default: ``1``. dtype (Dtype, optional): Data type of the output array. If unspecified the output type is inferred from the input array. Returns: array: Sum of specified diagonal. """ def transpose( a: array, /, axes: Sequence[int] | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Transpose the dimensions of the array. Args: a (array): Input array. axes (list(int), optional): Specifies the source axis for each axis in the new array. The default is to reverse the axes. Returns: array: The transposed array. """ def tri( n: int, m: int, k: int, dtype: Dtype | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ An array with ones at and below the given diagonal and zeros elsewhere. Args: n (int): The number of rows in the output. m (int, optional): The number of cols in the output. Defaults to ``None``. k (int, optional): The diagonal of the 2-D array. Defaults to ``0``. dtype (Dtype, optional): Data type of the output array. Defaults to ``float32``. stream (Stream, optional): Stream or device. Defaults to ``None``. Returns: array: Array with its lower triangle filled with ones and zeros elsewhere """ def tril(x: array, k: int, *, stream: Stream | Device | None = ...) -> array: """ Zeros the array above the given diagonal. Args: x (array): input array. k (int, optional): The diagonal of the 2-D array. Defaults to ``0``. stream (Stream, optional): Stream or device. Defaults to ``None``. Returns: array: Array zeroed above the given diagonal """ def triu(x: array, k: int, *, stream: Stream | Device | None = ...) -> array: """ Zeros the array below the given diagonal. Args: x (array): input array. k (int, optional): The diagonal of the 2-D array. Defaults to ``0``. stream (Stream, optional): Stream or device. Defaults to ``None``. Returns: array: Array zeroed below the given diagonal """ uint16: Dtype = ... uint32: Dtype = ... uint64: Dtype = ... uint8: Dtype = ... def unflatten( a: array, /, axis: int, shape: Sequence[int], *, stream: Stream | Device | None = ..., ) -> array: """ Unflatten an axis of an array to a shape. Args: a (array): Input array. axis (int): The axis to unflatten. shape (tuple(int)): The shape to unflatten to. At most one entry can be ``-1`` in which case the corresponding size will be inferred. stream (Stream, optional): Stream or device. Defaults to ``None`` in which case the default stream of the default device is used. Returns: array: The unflattened array. Example: >>> a = mx.array([1, 2, 3, 4]) >>> mx.unflatten(a, 0, (2, -1)) array([[1, 2], [3, 4]], dtype=int32) """ unsignedinteger: DtypeCategory = ... def value_and_grad( fun: Callable, argnums: int | Sequence[int] | None = ..., argnames: str | Sequence[str] = ..., ) -> Callable: """ Returns a function which computes the value and gradient of ``fun``. The function passed to :func:`value_and_grad` should return either a scalar loss or a tuple in which the first element is a scalar loss and the remaining elements can be anything. .. code-block:: python import mlx.core as mx def mse(params, inputs, targets): outputs = forward(params, inputs) lvalue = (outputs - targets).square().mean() return lvalue # Returns lvalue, dlvalue/dparams lvalue, grads = mx.value_and_grad(mse)(params, inputs, targets) def lasso(params, inputs, targets, a=1.0, b=1.0): outputs = forward(params, inputs) mse = (outputs - targets).square().mean() l1 = mx.abs(outputs - targets).mean() loss = a*mse + b*l1 return loss, mse, l1 (loss, mse, l1), grads = mx.value_and_grad(lasso)(params, inputs, targets) Args: fun (Callable): A function which takes a variable number of :class:`array` or trees of :class:`array` and returns a scalar output :class:`array` or a tuple the first element of which should be a scalar :class:`array`. argnums (int or list(int), optional): Specify the index (or indices) of the positional arguments of ``fun`` to compute the gradient with respect to. If neither ``argnums`` nor ``argnames`` are provided ``argnums`` defaults to ``0`` indicating ``fun``'s first argument. argnames (str or list(str), optional): Specify keyword arguments of ``fun`` to compute gradients with respect to. It defaults to [] so no gradients for keyword arguments by default. Returns: Callable: A function which returns a tuple where the first element is the output of `fun` and the second element is the gradients w.r.t. the loss. """ def var( a: array, /, axis: int | Sequence[int] | None = ..., keepdims: bool = ..., ddof: int = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Compute the variance(s) over the given axes. Args: a (array): Input array. axis (int or list(int), optional): Optional axis or axes to reduce over. If unspecified this defaults to reducing over the entire array. keepdims (bool, optional): Keep reduced axes as singleton dimensions, defaults to `False`. ddof (int, optional): The divisor to compute the variance is ``N - ddof``, defaults to 0. Returns: array: The output array of variances. """ def view( a: scalar | array, dtype: Dtype, stream: Stream | Device | None = ... ) -> array: """ View the array as a different type. The output shape changes along the last axis if the input array's type and the input ``dtype`` do not have the same size. Note: the view op does not imply that the input and output arrays share their underlying data. The view only gaurantees that the binary representation of each element (or group of elements) is the same. Args: a (array): Input array or scalar. dtype (Dtype): The data type to change to. Returns: array: The array with the new type. """ def vjp( fun: Callable, primals: list[array], cotangents: list[array] ) -> tuple[list[array], list[array]]: """ Compute the vector-Jacobian product. Computes the product of the ``cotangents`` with the Jacobian of a function ``fun`` evaluated at ``primals``. Args: fun (Callable): A function which takes a variable number of :class:`array` and returns a single :class:`array` or list of :class:`array`. primals (list(array)): A list of :class:`array` at which to evaluate the Jacobian. cotangents (list(array)): A list of :class:`array` which are the "vector" in the vector-Jacobian product. The ``cotangents`` should be the same in number, shape, and type as the outputs of ``fun``. Returns: list(array): A list of the vector-Jacobian products which is the same in number, shape, and type of the outputs of ``fun``. """ def vmap(fun: Callable, in_axes: object = ..., out_axes: object = ...) -> Callable: """ Returns a vectorized version of ``fun``. Args: fun (Callable): A function which takes a variable number of :class:`array` or a tree of :class:`array` and returns a variable number of :class:`array` or a tree of :class:`array`. in_axes (int, optional): An integer or a valid prefix tree of the inputs to ``fun`` where each node specifies the vmapped axis. If the value is ``None`` then the corresponding input(s) are not vmapped. Defaults to ``0``. out_axes (int, optional): An integer or a valid prefix tree of the outputs of ``fun`` where each node specifies the vmapped axis. If the value is ``None`` then the corresponding outputs(s) are not vmapped. Defaults to ``0``. Returns: Callable: The vectorized function. """ def where( condition: scalar | array, x: scalar | array, y: scalar | array, /, *, stream: Stream | Device | None = ..., ) -> array: """ Select from ``x`` or ``y`` according to ``condition``. The condition and input arrays must be the same shape or broadcastable with each another. Args: condition (array): The condition array. x (array): The input selected from where condition is ``True``. y (array): The input selected from where condition is ``False``. Returns: array: The output containing elements selected from ``x`` and ``y``. """ def zeros( shape: int | Sequence[int], dtype: Dtype | None = ..., *, stream: Stream | Device | None = ..., ) -> array: """ Construct an array of zeros. Args: shape (int or list(int)): The shape of the output array. dtype (Dtype, optional): Data type of the output array. If unspecified the output type defaults to ``float32``. Returns: array: The array of zeros with the specified shape. """ def zeros_like(a: array, /, *, stream: Stream | Device | None = ...) -> array: """ An array of zeros like the input. Args: a (array): The input to take the shape and type from. Returns: array: The output array filled with zeros. """ scalar: TypeAlias = int | float | bool list_or_scalar: TypeAlias = scalar | list["list_or_scalar"] bool_: Dtype = ... ================================================ FILE: .mlx_typings/mlx/core/cuda/__init__.pyi ================================================ def is_available() -> bool: """Check if the CUDA back-end is available.""" ================================================ FILE: .mlx_typings/mlx/core/distributed/__init__.pyi ================================================ from typing import Sequence from mlx.core import Device, Dtype, Stream, array class Group: """ An :class:`mlx.core.distributed.Group` represents a group of independent mlx processes that can communicate. """ def rank(self) -> int: """Get the rank of this process""" def size(self) -> int: """Get the size of the group""" def split(self, color: int, key: int = ...) -> Group: """ Split the group to subgroups based on the provided color. Processes that use the same color go to the same group. The ``key`` argument defines the rank in the new group. The smaller the key the smaller the rank. If the key is negative then the rank in the current group is used. Args: color (int): A value to group processes into subgroups. key (int, optional): A key to optionally change the rank ordering of the processes. """ def all_gather( x: array, *, group: Group | None = ..., stream: Stream | Device | None = ... ) -> array: """ Gather arrays from all processes. Gather the ``x`` arrays from all processes in the group and concatenate them along the first axis. The arrays should all have the same shape. Args: x (array): Input array. group (Group): The group of processes that will participate in the gather. If set to ``None`` the global group is used. Default: ``None``. stream (Stream, optional): Stream or device. Defaults to ``None`` in which case the default stream of the default device is used. Returns: array: The concatenation of all ``x`` arrays. """ def all_max( x: array, *, group: Group | None = ..., stream: Stream | Device | None = ... ) -> array: """ All reduce max. Find the maximum of the ``x`` arrays from all processes in the group. Args: x (array): Input array. group (Group): The group of processes that will participate in the reduction. If set to ``None`` the global group is used. Default: ``None``. stream (Stream, optional): Stream or device. Defaults to ``None`` in which case the default stream of the default device is used. Returns: array: The maximum of all ``x`` arrays. """ def all_min( x: array, *, group: Group | None = ..., stream: Stream | Device | None = ... ) -> array: """ All reduce min. Find the minimum of the ``x`` arrays from all processes in the group. Args: x (array): Input array. group (Group): The group of processes that will participate in the reduction. If set to ``None`` the global group is used. Default: ``None``. stream (Stream, optional): Stream or device. Defaults to ``None`` in which case the default stream of the default device is used. Returns: array: The minimum of all ``x`` arrays. """ def all_sum( x: array, *, group: Group | None = ..., stream: Stream | Device | None = ... ) -> array: """ All reduce sum. Sum the ``x`` arrays from all processes in the group. Args: x (array): Input array. group (Group): The group of processes that will participate in the reduction. If set to ``None`` the global group is used. Default: ``None``. stream (Stream, optional): Stream or device. Defaults to ``None`` in which case the default stream of the default device is used. Returns: array: The sum of all ``x`` arrays. """ def init(strict: bool = ..., backend: str = ...) -> Group: """ Initialize the communication backend and create the global communication group. Example: .. code:: python import mlx.core as mx group = mx.distributed.init(backend="ring") Args: strict (bool, optional): If set to False it returns a singleton group in case ``mx.distributed.is_available()`` returns False otherwise it throws a runtime error. Default: ``False`` backend (str, optional): Which distributed backend to initialize. Possible values ``mpi``, ``ring``, ``nccl``, ``any``. If set to ``any`` all available backends are tried and the first one that succeeds becomes the global group which will be returned in subsequent calls. Default: ``any`` Returns: Group: The group representing all the launched processes. """ def is_available() -> bool: """Check if a communication backend is available.""" def recv( shape: Sequence[int], dtype: Dtype, src: int, *, group: Group | None = ..., stream: Stream | Device | None = ..., ) -> array: """ Recv an array with shape ``shape`` and dtype ``dtype`` from process with rank ``src``. Args: shape (tuple[int]): The shape of the array we are receiving. dtype (Dtype): The data type of the array we are receiving. src (int): Rank of the source process in the group. group (Group): The group of processes that will participate in the recv. If set to ``None`` the global group is used. Default: ``None``. stream (Stream, optional): Stream or device. Defaults to ``None`` in which case the default stream of the default device is used. Returns: array: The array that was received from ``src``. """ def recv_like( x: array, src: int, *, group: Group | None = ..., stream: Stream | Device | None = ..., ) -> array: """ Recv an array with shape and type like ``x`` from process with rank ``src``. It is equivalent to calling ``mx.distributed.recv(x.shape, x.dtype, src)``. Args: x (array): An array defining the shape and dtype of the array we are receiving. src (int): Rank of the source process in the group. group (Group): The group of processes that will participate in the recv. If set to ``None`` the global group is used. Default: ``None``. stream (Stream, optional): Stream or device. Defaults to ``None`` in which case the default stream of the default device is used. Returns: array: The array that was received from ``src``. """ def send( x: array, dst: int, *, group: Group | None = ..., stream: Stream | Device | None = ..., ) -> array: """ Send an array from the current process to the process that has rank ``dst`` in the group. Args: x (array): Input array. dst (int): Rank of the destination process in the group. group (Group): The group of processes that will participate in the sned. If set to ``None`` the global group is used. Default: ``None``. stream (Stream, optional): Stream or device. Defaults to ``None`` in which case the default stream of the default device is used. Returns: array: An array identical to ``x`` which when evaluated the send is performed. """ ================================================ FILE: .mlx_typings/mlx/core/metal/__init__.pyi ================================================ def clear_cache() -> None: ... def device_info() -> dict[str, str | int]: """ Get information about the GPU device and system settings. Currently returns: * ``architecture`` * ``max_buffer_size`` * ``max_recommended_working_set_size`` * ``memory_size`` * ``resource_limit`` Returns: dict: A dictionary with string keys and string or integer values. """ def get_active_memory() -> int: ... def get_cache_memory() -> int: ... def get_peak_memory() -> int: ... def is_available() -> bool: """Check if the Metal back-end is available.""" def reset_peak_memory() -> None: ... def set_cache_limit(limit: int) -> int: ... def set_memory_limit(limit: int) -> int: ... def set_wired_limit(limit: int) -> int: ... def start_capture(path: str) -> None: """ Start a Metal capture. Args: path (str): The path to save the capture which should have the extension ``.gputrace``. """ def stop_capture() -> None: """Stop a Metal capture.""" ================================================ FILE: .mlx_typings/mlx/core/random/__init__.pyi ================================================ from typing import Sequence from mlx.core import Device, Dtype, Stream, array, scalar from mlx.core.distributed import state as state def bernoulli( p: scalar | array = ..., shape: Sequence[int] | None = ..., key: array | None = ..., stream: Stream | Device | None = ..., ) -> array: """ Generate Bernoulli random values. The values are sampled from the bernoulli distribution with parameter ``p``. The parameter ``p`` can be a :obj:`float` or :obj:`array` and must be broadcastable to ``shape``. Args: p (float or array, optional): Parameter of the Bernoulli distribution. Default: ``0.5``. shape (list(int), optional): Shape of the output. Default: ``p.shape``. key (array, optional): A PRNG key. Default: ``None``. Returns: array: The array of random integers. """ def categorical( logits: array, axis: int = ..., shape: Sequence[int] | None = ..., num_samples: int | None = ..., key: array | None = ..., stream: Stream | Device | None = ..., ) -> array: """ Sample from a categorical distribution. The values are sampled from the categorical distribution specified by the unnormalized values in ``logits``. Note, at most one of ``shape`` or ``num_samples`` can be specified. If both are ``None``, the output has the same shape as ``logits`` with the ``axis`` dimension removed. Args: logits (array): The *unnormalized* categorical distribution(s). axis (int, optional): The axis which specifies the distribution. Default: ``-1``. shape (list(int), optional): The shape of the output. This must be broadcast compatible with ``logits.shape`` with the ``axis`` dimension removed. Default: ``None`` num_samples (int, optional): The number of samples to draw from each of the categorical distributions in ``logits``. The output will have ``num_samples`` in the last dimension. Default: ``None``. key (array, optional): A PRNG key. Default: ``None``. Returns: array: The ``shape``-sized output array with type ``uint32``. """ def gumbel( shape: Sequence[int] = ..., dtype: Dtype | None = ..., key: Stream | Device | None = ..., stream: array | None = ..., ) -> array: """ Sample from the standard Gumbel distribution. The values are sampled from a standard Gumbel distribution which CDF ``exp(-exp(-x))``. Args: shape (list(int)): The shape of the output. dtype (Dtype, optional): The data type of the output. Default: ``float32``. key (array, optional): A PRNG key. Default: ``None``. Returns: array: The :class:`array` with shape ``shape`` and distributed according to the Gumbel distribution. """ def key(seed: int) -> array: """ Get a PRNG key from a seed. Args: seed (int): Seed for the PRNG. Returns: array: The PRNG key array. """ def laplace( shape: Sequence[int] = ..., dtype: Dtype | None = ..., loc: float = ..., scale: float = ..., key: array | None = ..., stream: Stream | Device | None = ..., ) -> array: """ Sample numbers from a Laplace distribution. Args: shape (list(int), optional): Shape of the output. Default: ``()``. dtype (Dtype, optional): Type of the output. Default: ``float32``. loc (float, optional): Mean of the distribution. Default: ``0.0``. scale (float, optional): The scale "b" of the Laplace distribution. Default:``1.0``. key (array, optional): A PRNG key. Default: ``None``. Returns: array: The output array of random values. """ def multivariate_normal( mean: array, cov: array, shape: Sequence[int] = ..., dtype: Dtype | None = ..., key: array | None = ..., stream: Stream | Device | None = ..., ) -> array: """ Generate jointly-normal random samples given a mean and covariance. The matrix ``cov`` must be positive semi-definite. The behavior is undefined if it is not. The only supported ``dtype`` is ``float32``. Args: mean (array): array of shape ``(..., n)``, the mean of the distribution. cov (array): array of shape ``(..., n, n)``, the covariance matrix of the distribution. The batch shape ``...`` must be broadcast-compatible with that of ``mean``. shape (list(int), optional): The output shape must be broadcast-compatible with ``mean.shape[:-1]`` and ``cov.shape[:-2]``. If empty, the result shape is determined by broadcasting the batch shapes of ``mean`` and ``cov``. Default: ``[]``. dtype (Dtype, optional): The output type. Default: ``float32``. key (array, optional): A PRNG key. Default: ``None``. Returns: array: The output array of random values. """ def normal( shape: Sequence[int] = ..., dtype: Dtype | None = ..., loc: scalar | array | None = ..., scale: scalar | array | None = ..., key: array | None = ..., stream: Stream | Device | None = ..., ) -> array: r""" Generate normally distributed random numbers. If ``loc`` and ``scale`` are not provided the "standard" normal distribution is used. That means $x \sim \mathcal{N}(0, 1)$ for real numbers and $\text{Re}(x),\text{Im}(x) \sim \mathcal{N}(0, \frac{1}{2})$ for complex numbers. Args: shape (list(int), optional): Shape of the output. Default: ``()``. dtype (Dtype, optional): Type of the output. Default: ``float32``. loc (scalar or array, optional): Mean of the distribution. Default: ``None``. scale (scalar or array, optional): Standard deviation of the distribution. Default: ``None``. key (array, optional): A PRNG key. Default: ``None``. Returns: array: The output array of random values. """ def permutation( x: int | array, axis: int = ..., key: array | None = ..., stream: Stream | Device | None = ..., ) -> array: """ Generate a random permutation or permute the entries of an array. Args: x (int or array, optional): If an integer is provided a random permtuation of ``mx.arange(x)`` is returned. Otherwise the entries of ``x`` along the given axis are randomly permuted. axis (int, optional): The axis to permute along. Default: ``0``. key (array, optional): A PRNG key. Default: ``None``. Returns: array: The generated random permutation or randomly permuted input array. """ def randint( low: scalar | array, high: scalar | array, shape: Sequence[int] = ..., dtype: Dtype | None = ..., key: array | None = ..., stream: Stream | Device | None = ..., ) -> array: """ Generate random integers from the given interval. The values are sampled with equal probability from the integers in half-open interval ``[low, high)``. The lower and upper bound can be scalars or arrays and must be broadcastable to ``shape``. Args: low (scalar or array): Lower bound of the interval. high (scalar or array): Upper bound of the interval. shape (list(int), optional): Shape of the output. Default: ``()``. dtype (Dtype, optional): Type of the output. Default: ``int32``. key (array, optional): A PRNG key. Default: ``None``. Returns: array: The array of random integers. """ def seed(seed: int) -> None: """ Seed the global PRNG. Args: seed (int): Seed for the global PRNG. """ def split(key: array, num: int = ..., stream: Stream | Device | None = ...) -> array: """ Split a PRNG key into sub keys. Args: key (array): Input key to split. num (int, optional): Number of sub keys. Default: ``2``. Returns: array: The array of sub keys with ``num`` as its first dimension. """ def truncated_normal( lower: scalar | array, upper: scalar | array, shape: Sequence[int] | None = ..., dtype: Dtype | None = ..., key: array | None = ..., stream: Stream | Device | None = ..., ) -> array: """ Generate values from a truncated normal distribution. The values are sampled from the truncated normal distribution on the domain ``(lower, upper)``. The bounds ``lower`` and ``upper`` can be scalars or arrays and must be broadcastable to ``shape``. Args: lower (scalar or array): Lower bound of the domain. upper (scalar or array): Upper bound of the domain. shape (list(int), optional): The shape of the output. Default:``()``. dtype (Dtype, optional): The data type of the output. Default: ``float32``. key (array, optional): A PRNG key. Default: ``None``. Returns: array: The output array of random values. """ def uniform( low: scalar | array = ..., high: scalar | array = ..., shape: Sequence[int] = ..., dtype: Dtype | None = ..., key: array | None = ..., stream: Stream | Device | None = ..., ) -> array: """ Generate uniformly distributed random numbers. The values are sampled uniformly in the half-open interval ``[low, high)``. The lower and upper bound can be scalars or arrays and must be broadcastable to ``shape``. Args: low (scalar or array, optional): Lower bound of the distribution. Default: ``0``. high (scalar or array, optional): Upper bound of the distribution. Default: ``1``. shape (list(int), optional): Shape of the output. Default:``()``. dtype (Dtype, optional): Type of the output. Default: ``float32``. key (array, optional): A PRNG key. Default: ``None``. Returns: array: The output array random values. """ ================================================ FILE: .mlx_typings/mlx/nn/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ from layers import * from utils import * from . import init as init from . import losses as losses ================================================ FILE: .mlx_typings/mlx/nn/init.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Callable, Literal import mlx.core as mx def constant(value: float, dtype: mx.Dtype = ...) -> Callable[[mx.array], mx.array]: r"""An initializer that returns an array filled with ``value``. Args: value (float): The value to fill the array with. dtype (Dtype, optional): The data type of the array. Default: ``float32``. Returns: Callable[[array], array]: An initializer that returns an array with the same shape as the input, filled with ``value``. Example: >>> init_fn = nn.init.constant(0.5) >>> init_fn(mx.zeros((2, 2))) array([[0.5, 0.5], [0.5, 0.5]], dtype=float32) """ def normal( mean: float = ..., std: float = ..., dtype: mx.Dtype = ... ) -> Callable[[mx.array], mx.array]: r"""An initializer that returns samples from a normal distribution. Args: mean (float, optional): Mean of the normal distribution. Default: ``0.0``. std (float, optional): Standard deviation of the normal distribution. Default: ``1.0``. dtype (Dtype, optional): The data type of the array. Default: ``float32``. Returns: Callable[[array], array]: An initializer that returns an array with the same shape as the input, filled with samples from a normal distribution. Example: >>> init_fn = nn.init.normal() >>> init_fn(mx.zeros((2, 2))) array([[-0.982273, -0.534422], [0.380709, 0.0645099]], dtype=float32) """ def uniform( low: float = ..., high: float = ..., dtype: mx.Dtype = ... ) -> Callable[[mx.array], mx.array]: r"""An initializer that returns samples from a uniform distribution. Args: low (float, optional): The lower bound of the uniform distribution. Default: ``0.0``. high (float, optional): The upper bound of the uniform distribution. Default: ``1.0`` dtype (Dtype, optional): The data type of the array. Default: ``float32``. Returns: Callable[[array], array]: An initializer that returns an array with the same shape as the input, filled with samples from a uniform distribution Example: >>> init_fn = nn.init.uniform(low=0, high=1) >>> init_fn(mx.zeros((2, 2))) array([[0.883935, 0.863726], [0.617261, 0.417497]], dtype=float32) """ def identity(dtype: mx.Dtype = ...) -> Callable[[mx.array], mx.array]: r"""An initializer that returns an identity matrix. Args: dtype (Dtype, optional): The data type of the array. Defaults: ``float32``. Returns: Callable[[array], array]: An initializer that returns an identity matrix with the same shape as the input. Example: >>> init_fn = nn.init.identity() >>> init_fn(mx.zeros((2, 2))) array([[1, 0], [0, 1]], dtype=float32) """ def glorot_normal(dtype: mx.Dtype = ...) -> Callable[[mx.array, float], mx.array]: r"""A Glorot normal initializer. This initializer samples from a normal distribution with a standard deviation computed from the number of input (``fan_in``) and output (``fan_out``) units according to: .. math:: \sigma = \gamma \sqrt{\frac{2.0}{\text{fan\_in} + \text{fan\_out}}} For more details see the original reference: `Understanding the difficulty of training deep feedforward neural networks `_ Args: dtype (Dtype, optional): The data type of the array. Default: ``float32``. Returns: Callable[[array, float], array]: An initializer that returns an array with the same shape as the input, filled with samples from the Glorot normal distribution. Example: >>> init_fn = nn.init.glorot_normal() >>> init_fn(mx.zeros((2, 2))) array([[0.191107, 1.61278], [-0.150594, -0.363207]], dtype=float32) >>> init_fn(mx.zeros((2, 2)), gain=4.0) array([[1.89613, -4.53947], [4.48095, 0.995016]], dtype=float32) """ def glorot_uniform(dtype: mx.Dtype = ...) -> Callable[[mx.array, float], mx.array]: r"""A Glorot uniform initializer. This initializer samples from a uniform distribution with a range computed from the number of input (``fan_in``) and output (``fan_out``) units according to: .. math:: \sigma = \gamma \sqrt{\frac{6.0}{\text{fan\_in} + \text{fan\_out}}} For more details see the original reference: `Understanding the difficulty of training deep feedforward neural networks `_ Args: dtype (Dtype, optional): The data type of the array. Default: ``float32``. Returns: Callable[[array, float], array]: An initializer that returns an array with the same shape as the input, filled with samples from the Glorot uniform distribution. Example: >>> init_fn = nn.init.glorot_uniform() >>> init_fn(mx.zeros((2, 2))) array([[0.223404, -0.890597], [-0.379159, -0.776856]], dtype=float32) >>> init_fn(mx.zeros((2, 2)), gain=4.0) array([[-1.90041, 3.02264], [-0.912766, 4.12451]], dtype=float32) """ def he_normal( dtype: mx.Dtype = ..., ) -> Callable[[mx.array, Literal["fan_in", "fan_out"], float], mx.array]: r"""Build a He normal initializer. This initializer samples from a normal distribution with a standard deviation computed from the number of input (``fan_in``) or output (``fan_out``) units according to: .. math:: \sigma = \gamma \frac{1}{\sqrt{\text{fan}}} where :math:`\text{fan}` is either the number of input units when the ``mode`` is ``"fan_in"`` or output units when the ``mode`` is ``"fan_out"``. For more details see the original reference: `Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification `_ Args: dtype (Dtype, optional): The data type of the array. Defaults to mx.float32. Returns: Callable[[array, str, float], array]: An initializer that returns an array with the same shape as the input, filled with samples from the He normal distribution. Example: >>> init_fn = nn.init.he_normal() >>> init_fn(mx.zeros((2, 2))) # uses fan_in array([[-1.25211, 0.458835], [-0.177208, -0.0137595]], dtype=float32) >>> init_fn(mx.zeros((2, 2)), mode="fan_out", gain=5) array([[5.6967, 4.02765], [-4.15268, -2.75787]], dtype=float32) """ def he_uniform( dtype: mx.Dtype = ..., ) -> Callable[[mx.array, Literal["fan_in", "fan_out"], float], mx.array]: r"""A He uniform (Kaiming uniform) initializer. This initializer samples from a uniform distribution with a range computed from the number of input (``fan_in``) or output (``fan_out``) units according to: .. math:: \sigma = \gamma \sqrt{\frac{3.0}{\text{fan}}} where :math:`\text{fan}` is either the number of input units when the ``mode`` is ``"fan_in"`` or output units when the ``mode`` is ``"fan_out"``. For more details see the original reference: `Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification `_ Args: dtype (Dtype, optional): The data type of the array. Default: ``float32``. Returns: Callable[[array, str, float], array]: An initializer that returns an array with the same shape as the input, filled with samples from the He uniform distribution. Example: >>> init_fn = nn.init.he_uniform() >>> init_fn(mx.zeros((2, 2))) # uses fan_in array([[0.0300242, -0.0184009], [0.793615, 0.666329]], dtype=float32) >>> init_fn(mx.zeros((2, 2)), mode="fan_out", gain=5) array([[-1.64331, -2.16506], [1.08619, 5.79854]], dtype=float32) """ def sparse( sparsity: float, mean: float = ..., std: float = ..., dtype: mx.Dtype = ... ) -> Callable[[mx.array], mx.array]: r"""An initializer that returns a sparse matrix. Args: sparsity (float): The fraction of elements in each column to be set to zero. mean (float, optional): Mean of the normal distribution. Default: ``0.0``. std (float, optional): Standard deviation of the normal distribution. Default: ``1.0``. dtype (Dtype, optional): The data type of the array. Default: ``float32``. Returns: Callable[[array], array]: An initializer that returns an array with the same shape as the input, filled with samples from a normal distribution. Example: >>> init_fn = nn.init.sparse(sparsity=0.5) >>> init_fn(mx.zeros((2, 2))) array([[-1.91187, -0.117483], [0, 0]], dtype=float32) """ def orthogonal( gain: float = ..., dtype: mx.Dtype = ... ) -> Callable[[mx.array], mx.array]: r"""An initializer that returns an orthogonal matrix. Args: gain (float, optional): Scaling factor for the orthogonal matrix. Default: ``1.0``. dtype (Dtype, optional): Data type of the array. Default: ``float32``. Returns: Callable[[array], array]: An initializer that returns an orthogonal matrix with the same shape as the input. """ ================================================ FILE: .mlx_typings/mlx/nn/layers/__init__.pyi ================================================ """ This type stub file was generated by pyright. """ from activations import * from base import * from containers import * from convolution import * from convolution_transpose import * from distributed import * from dropout import * from embedding import * from linear import * from normalization import * from pooling import * from positional_encoding import * from quantized import * from recurrent import * from transformer import * from upsample import * ================================================ FILE: .mlx_typings/mlx/nn/layers/activations.pyi ================================================ """ This type stub file was generated by pyright. """ from functools import partial from typing import Any import mlx.core as mx from base import Module @partial(mx.compile, shapeless=True) def sigmoid(x: mx.array) -> mx.array: r"""Applies the sigmoid function. .. math:: \text{Sigmoid}(x) = \sigma(x) = \frac{1}{1 + \exp(-x)} """ @partial(mx.compile, shapeless=True) def relu(x: mx.array) -> mx.array: r"""Applies the Rectified Linear Unit. Simply ``mx.maximum(x, 0)``. """ @partial(mx.compile, shapeless=True) def relu2(x: mx.array) -> mx.array: r"""Applies the ReLU² activation function. Applies :math:`\max(0, x)^2` element wise. """ @partial(mx.compile, shapeless=True) def relu6(x: mx.array) -> mx.array: r"""Applies the Rectified Linear Unit 6. Applies :math:`\min(\max(x, 0), 6)` element wise. """ @partial(mx.compile, shapeless=True) def leaky_relu(x: mx.array, negative_slope=...) -> mx.array: r"""Applies the Leaky Rectified Linear Unit. Simply ``mx.maximum(negative_slope * x, x)``. """ @partial(mx.compile, shapeless=True) def log_softmax(x: mx.array, axis=...): r"""Applies the Log Softmax function. Applies :math:`x + \log \sum_i e^{x_i}` element wise. """ @partial(mx.compile, shapeless=True) def elu(x: mx.array, alpha=...) -> mx.array: r"""Applies the Exponential Linear Unit. Simply ``mx.where(x > 0, x, alpha * (mx.exp(x) - 1))``. """ @partial(mx.compile, shapeless=True) def softmax(x: mx.array, axis=...) -> mx.array: r"""Applies the Softmax function. Applies :math:`\frac{e^{x_i}}{\sum_j e^{x_j}}` element wise. """ @partial(mx.compile, shapeless=True) def softplus(x: mx.array) -> mx.array: r"""Applies the Softplus function. Applies :math:`\log(1 + \exp(x))` element wise. """ @partial(mx.compile, shapeless=True) def softsign(x: mx.array) -> mx.array: r"""Applies the Softsign function. Applies :math:`\frac{x}{1 + |x|}` element wise. """ @partial(mx.compile, shapeless=True) def softshrink(x: mx.array, lambd: float = ...) -> mx.array: r"""Applies the Softshrink activation function. .. math:: \text{softshrink}(x) = \begin{cases} x - \lambda & \text{if } x > \lambda \\ x + \lambda & \text{if } x < -\lambda \\ 0 & \text{otherwise} \end{cases} """ @partial(mx.compile, shapeless=True) def celu(x: mx.array, alpha=...) -> mx.array: r"""Applies the Continuously Differentiable Exponential Linear Unit. Applies :math:`\max(0, x) + \min(0, \alpha * (\exp(x / \alpha) - 1))` element wise. """ @partial(mx.compile, shapeless=True) def silu(x: mx.array) -> mx.array: r"""Applies the Sigmoid Linear Unit. Also known as Swish. Applies :math:`x \sigma(x)` element wise, where :math:`\sigma(\cdot)` is the logistic sigmoid. """ @partial(mx.compile, shapeless=True) def log_sigmoid(x: mx.array) -> mx.array: r"""Applies the Log Sigmoid function. Applies :math:`\log(\sigma(x)) = -\log(1 + e^{-x})` element wise. """ @partial(mx.compile, shapeless=True) def gelu(x: mx.array) -> mx.array: r"""Applies the Gaussian Error Linear Units function. .. math:: \textrm{GELU}(x) = x * \Phi(x) where :math:`\Phi(x)` is the Gaussian CDF. See also :func:`gelu_approx` and :func:`gelu_fast_approx` for faster approximations. """ @partial(mx.compile, shapeless=True) def gelu_approx(x: mx.array) -> mx.array: r"""An approximation to Gaussian Error Linear Unit. See :func:`gelu` for the exact computation. This function approximates ``gelu`` with a maximum absolute error :math:`< 0.0005` in the range :math:`[-6, 6]` using the following .. math:: x = 0.5 * x * \left(1 + \text{Tanh}\left((\sqrt{2 / \pi} * \left(x + 0.044715 * x^3\right)\right)\right) """ @partial(mx.compile, shapeless=True) def gelu_fast_approx(x: mx.array) -> mx.array: r"""A fast approximation to Gaussian Error Linear Unit. See :func:`gelu` for the exact computation. This function approximates ``gelu`` with a maximum absolute error :math:`< 0.015` in the range :math:`[-6, 6]` using the following .. math:: x = x \sigma\left(1.702 x\right) where :math:`\sigma(\cdot)` is the logistic sigmoid. References: - https://github.com/hendrycks/GELUs - https://arxiv.org/abs/1606.08415 """ def glu(x: mx.array, axis: int = ...) -> mx.array: r"""Applies the gated linear unit function. This function splits the ``axis`` dimension of the input into two halves (:math:`a` and :math:`b`) and applies :math:`a * \sigma(b)`. .. math:: \textrm{GLU}(x) = a * \sigma(b) Args: axis (int): The dimension to split along. Default: ``-1`` """ @partial(mx.compile, shapeless=True) def step(x: mx.array, threshold: float = ...) -> mx.array: r"""Applies the Step Activation Function. This function implements a binary step activation, where the output is set to 1 if the input is greater than a specified threshold, and 0 otherwise. .. math:: \text{step}(x) = \begin{cases} 0 & \text{if } x < \text{threshold} \\ 1 & \text{if } x \geq \text{threshold} \end{cases} Args: threshold: The value to threshold at. """ @partial(mx.compile, shapeless=True) def selu(x: mx.array) -> mx.array: r"""Applies the Scaled Exponential Linear Unit. .. math:: \text{selu}(x) = \begin{cases} \lambda x & \text{if } x > 0 \\ \lambda \alpha (\exp(x) - 1) & \text{if } x \leq 0 \end{cases} where :math:`\lambda = 1.0507` and :math:`\alpha = 1.67326`. See also :func:`elu`. """ @partial(mx.compile, shapeless=True) def prelu(x: mx.array, alpha: mx.array) -> mx.array: r"""Applies the element-wise parametric ReLU. .. math:: \text{PReLU}(x) = \max(0,x) + a * \min(0,x) where :math:`a` is an array. """ @partial(mx.compile, shapeless=True) def mish(x: mx.array) -> mx.array: r"""Applies the Mish function, element-wise. Mish: A Self Regularized Non-Monotonic Neural Activation Function. Reference: https://arxiv.org/abs/1908.08681 .. math:: \text{Mish}(x) = x * \text{Tanh}(\text{Softplus}(x)) """ @partial(mx.compile, shapeless=True) def hardswish(x: mx.array) -> mx.array: r"""Applies the hardswish function, element-wise. .. math:: \text{Hardswish}(x) = x * \min(\max(x + 3, 0), 6) / 6 """ @partial(mx.compile, shapeless=True) def hard_tanh(x: mx.array, min_val=..., max_val=...) -> mx.array: r"""Applies the HardTanh function. Applies :math:`\max(\min(x, \text{max\_val}), \text{min\_val})` element-wise. """ @partial(mx.compile, shapeless=True) def hard_shrink(x: mx.array, lambd=...) -> mx.array: r"""Applies the HardShrink activation function. .. math:: \text{hardshrink}(x) = \begin{cases} x & \text{if } x > \lambda \\ x & \text{if } x < -\lambda \\ 0 & \text{otherwise} \end{cases} """ @partial(mx.compile, shapeless=True) def softmin(x: mx.array, axis=...) -> mx.array: r"""Applies the Softmin function. Applies :math:`\frac{e^{-x_i}}{\sum_j e^{-x_j}}` element-wise. """ def tanh(x: mx.array) -> mx.array: """Applies the hyperbolic tangent function. Simply ``mx.tanh(x)``. """ class GLU(Module): r"""Applies the gated linear unit function. This function splits the ``axis`` dimension of the input into two halves (:math:`a` and :math:`b`) and applies :math:`a * \sigma(b)`. .. math:: \textrm{GLU}(x) = a * \sigma(b) Args: axis (int): The dimension to split along. Default: ``-1`` """ def __init__(self, axis: int = ...) -> None: ... def __call__(self, x) -> Any: ... @_make_activation_module(sigmoid) class Sigmoid(Module): r"""Applies the sigmoid function, element-wise. .. math:: \text{Sigmoid}(x) = \sigma(x) = \frac{1}{1 + \exp(-x)} """ @_make_activation_module(mish) class Mish(Module): r"""Applies the Mish function, element-wise. Reference: https://arxiv.org/abs/1908.08681 .. math:: \text{Mish}(x) = x * \text{Tanh}(\text{Softplus}(x)) """ @_make_activation_module(relu) class ReLU(Module): r"""Applies the Rectified Linear Unit. Simply ``mx.maximum(x, 0)``. See :func:`relu` for the functional equivalent. """ @_make_activation_module(relu2) class ReLU2(Module): r"""Applies the ReLU² activation function. See :func:`relu2` for the functional equivalent. """ @_make_activation_module(relu6) class ReLU6(Module): r"""Applies the Rectified Linear Unit 6. See :func:`relu6` for the functional equivalent. """ class LeakyReLU(Module): r"""Applies the Leaky Rectified Linear Unit. Simply ``mx.maximum(negative_slope * x, x)``. Args: negative_slope: Controls the angle of the negative slope. Default: ``1e-2`` """ def __init__(self, negative_slope=...) -> None: ... def __call__(self, x): ... class ELU(Module): r"""Applies the Exponential Linear Unit. Simply ``mx.where(x > 0, x, alpha * (mx.exp(x) - 1))``. See :func:`elu` for the functional equivalent. Args: alpha: the :math:`\alpha` value for the ELU formulation. Default: ``1.0`` """ def __init__(self, alpha=...) -> None: ... def __call__(self, x): ... @_make_activation_module(softmax) class Softmax(Module): r"""Applies the Softmax function. See :func:`softmax` for the functional equivalent. """ @_make_activation_module(softplus) class Softplus(Module): r"""Applies the Softplus function. See :func:`softplus` for the functional equivalent. """ @_make_activation_module(softsign) class Softsign(Module): r"""Applies the Softsign function. See :func:`softsign` for the functional equivalent. """ class Softshrink(Module): r"""Applies the Softshrink function. See :func:`softshrink` for the functional equivalent. Args: lambd: the :math:`\lambda` value for Softshrink. Default: ``0.5`` """ def __init__(self, lambd=...) -> None: ... def __call__(self, x): ... class CELU(Module): r"""Applies the Continuously Differentiable Exponential Linear Unit. Applies :math:`\max(0, x) + \min(0, \alpha * (\exp(x / \alpha) - 1))` element wise. See :func:`celu` for the functional equivalent. Args: alpha: the :math:`\alpha` value for the CELU formulation. Default: ``1.0`` """ def __init__(self, alpha=...) -> None: ... def __call__(self, x): ... @_make_activation_module(silu) class SiLU(Module): r"""Applies the Sigmoid Linear Unit. Also known as Swish. See :func:`silu` for the functional equivalent. """ @_make_activation_module(log_softmax) class LogSoftmax(Module): r"""Applies the Log Softmax function. See :func:`log_softmax` for the functional equivalent. """ @_make_activation_module(log_sigmoid) class LogSigmoid(Module): r"""Applies the Log Sigmoid function. See :func:`log_sigmoid` for the functional equivalent. """ class PReLU(Module): r"""Applies the element-wise parametric ReLU. Applies :math:`\max(0, x) + a * \min(0, x)` element wise, where :math:`a` is an array. See :func:`prelu` for the functional equivalent. Args: num_parameters: number of :math:`a` to learn. Default: ``1`` init: the initial value of :math:`a`. Default: ``0.25`` """ def __init__(self, num_parameters=..., init=...) -> None: ... def __call__(self, x: mx.array): ... class GELU(Module): r"""Applies the Gaussian Error Linear Units. .. math:: \textrm{GELU}(x) = x * \Phi(x) where :math:`\Phi(x)` is the Gaussian CDF. However, if ``approx`` is set to 'precise' or 'fast' it applies .. math:: \textrm{GELUApprox}(x) &= 0.5 * x * \left(1 + \text{Tanh}\left((\sqrt{2 / \pi} * \left(x + 0.044715 * x^3\right)\right)\right) \\ \textrm{GELUFast}(x) &= x * \sigma\left(1.702 * x\right) respectively. .. note:: For compatibility with the PyTorch API, 'tanh' can be used as an alias for 'precise'. See :func:`gelu`, :func:`gelu_approx` and :func:`gelu_fast_approx` for the functional equivalents and information regarding error bounds. Args: approx ('none' | 'precise' | 'fast'): Which approximation to gelu to use if any. """ def __init__(self, approx=...) -> None: ... def __call__(self, x): ... @_make_activation_module(tanh) class Tanh(Module): r"""Applies the hyperbolic tangent function. See :func:`tanh` for the functional equivalent. """ @_make_activation_module(hardswish) class Hardswish(Module): r"""Applies the hardswish function, element-wise. See :func:`hardswish` for the functional equivalent. """ class Step(Module): r"""Applies the Step Activation Function. This function implements a binary step activation, where the output is set to 1 if the input is greater than a specified threshold, and 0 otherwise. .. math:: \text{step}(x) = \begin{cases} 0 & \text{if } x < \text{threshold} \\ 1 & \text{if } x \geq \text{threshold} \end{cases} Args: threshold: The value to threshold at. """ def __init__(self, threshold: float = ...) -> None: ... def __call__(self, x: mx.array): ... @_make_activation_module(selu) class SELU(Module): r"""Applies the Scaled Exponential Linear Unit. See :func:`selu` for the functional equivalent. """ @_make_activation_module(hard_tanh) class HardTanh(Module): r"""Applies the HardTanh function. See :func:`hard_tanh` for the functional equivalent. """ @_make_activation_module(hard_shrink) class HardShrink(Module): r"""Applies the HardShrink function. See :func:`hard_shrink` for the functional equivalent. Args: lambd: the :math:`\lambda` value for Hardshrink. Default: ``0.5`` """ @_make_activation_module(softmin) class Softmin(Module): r"""Applies the Softmin function. See :func:`softmin` for the functional equivalent. """ ================================================ FILE: .mlx_typings/mlx/nn/layers/base.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Any, Callable, List, Optional, Tuple, Union import mlx.core as mx class Module(dict): """Base class for building neural networks with MLX. All the layers provided in :mod:`layers` subclass this class and your models should do the same. A ``Module`` can contain other ``Module`` instances or :class:`mlx.core.array` instances in arbitrary nesting of python lists or dicts. The ``Module`` then allows recursively extracting all the :class:`mlx.core.array` instances using :meth:`Module.parameters`. In addition, the ``Module`` has the concept of trainable and non trainable parameters (called "frozen"). When using :func:`value_and_grad` the gradients are returned only with respect to the trainable parameters. All arrays in a module are trainable unless they are added in the "frozen" set by calling :meth:`freeze`. .. code-block:: python import mlx.core as mx import mlx.nn as nn class MyMLP(nn.Module): def __init__(self, in_dims: int, out_dims: int, hidden_dims: int = 16): super().__init__() self.in_proj = nn.Linear(in_dims, hidden_dims) self.out_proj = nn.Linear(hidden_dims, out_dims) def __call__(self, x): x = self.in_proj(x) x = mx.maximum(x, 0) return self.out_proj(x) model = MyMLP(2, 1) # All the model parameters are created but since MLX is lazy by # default, they are not evaluated yet. Calling `mx.eval` actually # allocates memory and initializes the parameters. mx.eval(model.parameters()) # Setting a parameter to a new value is as simply as accessing that # parameter and assigning a new array to it. model.in_proj.weight = model.in_proj.weight * 2 mx.eval(model.parameters()) """ __call__: Callable def __init__(self) -> None: """Should be called by the subclasses of ``Module``.""" @property def training(self): # -> bool: """Boolean indicating if the model is in training mode.""" @property def state(self): # -> Self: """The module's state dictionary The module's state dictionary contains any attribute set on the module including parameters in :meth:`Module.parameters` Unlike :meth:`Module.parameters`, the :attr:`Module.state` property is a reference to the module's state. Updates to it will be reflected in the original module. """ def __repr__(self): # -> str: ... def __getattr__(self, key: str): # -> None: ... def __setattr__(self, key: str, val: Any): # -> None: ... def __delattr__(self, name): # -> None: ... def load_weights( self, file_or_weights: Union[str, List[Tuple[str, mx.array]]], strict: bool = ..., ) -> Module: """ Update the model's weights from a ``.npz``, a ``.safetensors`` file, or a list. Args: file_or_weights (str or list(tuple(str, mx.array))): The path to the weights ``.npz`` file (``.npz`` or ``.safetensors``) or a list of pairs of parameter names and arrays. strict (bool, optional): If ``True`` then checks that the provided weights exactly match the parameters of the model. Otherwise, only the weights actually contained in the model are loaded and shapes are not checked. Default: ``True``. Returns: The module instance after updating the weights. Example: .. code-block:: python import mlx.core as mx import mlx.nn as nn model = nn.Linear(10, 10) # Load from file model.load_weights("weights.npz") # Load from .safetensors file model.load_weights("weights.safetensors") # Load from list weights = [ ("weight", mx.random.uniform(shape=(10, 10))), ("bias", mx.zeros((10,))), ] model.load_weights(weights) # Missing weight weights = [ ("weight", mx.random.uniform(shape=(10, 10))), ] # Raises a ValueError exception model.load_weights(weights) # Ok, only updates the weight but not the bias model.load_weights(weights, strict=False) """ def save_weights(self, file: str): # -> None: """ Save the model's weights to a file. The saving method is determined by the file extension: - ``.npz`` will use :func:`mx.savez` - ``.safetensors`` will use :func:`mx.save_safetensors` """ @staticmethod def is_module(value): # -> bool: ... @staticmethod def valid_child_filter(module, key, value): # -> bool: ... @staticmethod def valid_parameter_filter(module, key, value): # -> bool: ... @staticmethod def trainable_parameter_filter(module, key, value): # -> bool: ... def filter_and_map( self, filter_fn: Callable[[Module, str, Any], bool], map_fn: Optional[Callable] = ..., is_leaf_fn: Optional[Callable[[Module, str, Any], bool]] = ..., ): # -> dict[Any, Any | dict[Any, Any | dict[Any, Any] | list[Any]] | dict[Any, Any] | list[Any]]: """Recursively filter the contents of the module using ``filter_fn``, namely only select keys and values where ``filter_fn`` returns true. This is used to implement :meth:`parameters` and :meth:`trainable_parameters` but it can also be used to extract any subset of the module's parameters. Args: filter_fn (Callable): Given a value, the key in which it is found and the containing module, decide whether to keep the value or drop it. map_fn (Callable, optional): Optionally transform the value before returning it. is_leaf_fn (Callable, optional): Given a value, the key in which it is found and the containing module decide if it is a leaf. Returns: A dictionary containing the contents of the module recursively filtered """ def parameters( self, ) -> mx.MX_ARRAY_TREE: """Recursively return all the :class:`mlx.core.array` members of this Module as a dict of dicts and lists.""" def trainable_parameters( self, ) -> mx.MX_ARRAY_TREE: # -> dict[Any, Any | dict[Any, Any | dict[Any, Any] | list[Any]] | dict[Any, Any] | list[Any]]: """Recursively return all the non frozen :class:`mlx.core.array` members of this Module as a dict of dicts and lists.""" def children( self, ) -> mx.MX_ARRAY_TREE: # -> dict[Any, Any | dict[Any, Any | dict[Any, Any] | list[Any]] | dict[Any, Any] | list[Any]]: """Return the direct descendants of this Module instance.""" def leaf_modules( self, ) -> mx.MX_ARRAY_TREE: # -> dict[Any, Any | dict[Any, Any | dict[Any, Any] | list[Any]] | dict[Any, Any] | list[Any]]: """Return the submodules that do not contain other modules.""" def update(self, parameters: dict[str, Any], strict: bool = ...) -> Module: """Replace the parameters of this Module with the provided ones in the dict of dicts and lists. Commonly used by the optimizer to change the model to the updated (optimized) parameters. Also used by the :meth:`value_and_grad` to set the tracers in the model in order to compute gradients. The passed in parameters dictionary need not be a full dictionary similar to :meth:`parameters`. Only the provided locations will be updated. Args: parameters (dict): A complete or partial dictionary of the modules parameters. strict (bool): If ``True`` checks that ``parameters`` is a subset of the module's parameters. Default: ``True``. Returns: The module instance after updating the parameters. """ def apply( self, map_fn: Callable[[mx.array], mx.array], filter_fn: Optional[Callable[[Module, str, Any], bool]] = ..., ) -> Module: """Map all the parameters using the provided ``map_fn`` and immediately update the module with the mapped parameters. For instance running ``model.apply(lambda x: x.astype(mx.float16))`` casts all parameters to 16 bit floats. Args: map_fn (Callable): Maps an array to another array filter_fn (Callable, optional): Filter to select which arrays to map (default: :meth:`Module.valid_parameter_filter`). Returns: The module instance after updating the parameters. """ def update_modules(self, modules: dict, strict: bool = ...) -> Module: """Replace the child modules of this :class:`Module` instance with the provided ones in the dict of dicts and lists. It is the equivalent of :meth:`Module.update` but for modules instead of parameters and allows us to flexibly edit complex architectures by programmatically swapping layers. The passed in parameters dictionary need not be a full dictionary similar to :meth:`modules`. Only the provided locations will be updated. Args: modules (dict): A complete or partial dictionary of the module's submodules. strict (bool): If ``True`` checks that ``modules`` is a subset of the child modules of this instance. Default: ``True``. Returns: The module instance after updating the submodules. """ def apply_to_modules(self, apply_fn: Callable[[str, Module], Any]) -> Module: """Apply a function to all the modules in this instance (including this instance). Args: apply_fn (Callable): The function to apply to the modules. Returns: The module instance after updating submodules. """ def modules(self): # -> list[Any]: """Return a list with all the modules in this instance. Returns: A list of :class:`Module` instances. """ def named_modules(self): # -> list[Any]: """Return a list with all the modules in this instance and their name with dot notation. Returns: A list of tuples (str, :class:`Module`). """ def freeze( self, *, recurse: bool = ..., keys: Optional[Union[str, List[str]]] = ..., strict: bool = ..., ) -> Module: """Freeze the Module's parameters or some of them. Freezing a parameter means not computing gradients for it. This function is idempotent i.e. freezing a frozen model is a no-op. Example: For instance to only train the attention parameters from a Transformer: .. code-block:: python model = nn.Transformer() model.freeze() model.apply_to_modules(lambda k, v: v.unfreeze() if k.endswith("attention") else None) Args: recurse (bool, optional): If True then freeze the parameters of the submodules as well. Default: ``True``. keys (str or list[str], optional): If provided then only these parameters will be frozen otherwise all the parameters of a module. For instance freeze all biases by calling ``module.freeze(keys="bias")``. strict (bool, optional): If set to ``True`` validate that the passed keys exist. Default: ``False``. Returns: The module instance after freezing the parameters. """ def unfreeze( self, *, recurse: bool = ..., keys: Optional[Union[str, List[str]]] = ..., strict: bool = ..., ) -> Module: """Unfreeze the Module's parameters or some of them. This function is idempotent ie unfreezing a model that is not frozen is a noop. Example: For instance to only train the biases of a Transformer one can do: .. code-block:: python model = nn.Transformer() model.freeze() model.unfreeze(keys="bias") Args: recurse (bool, optional): If True then unfreeze the parameters of the submodules as well. Default: ``True``. keys (str or list[str], optional): If provided then only these parameters will be unfrozen otherwise all the parameters of a module. For instance unfreeze all biases by calling ``module.unfreeze(keys="bias")``. strict (bool, optional): If set to ``True`` validate that the passed keys exist. Default: ``False``. Returns: The module instance after unfreezing the parameters. """ def train(self, mode: bool = ...) -> Module: """Set the model in or out of training mode. Training mode only applies to certain layers. For example :obj:`Dropout` applies a random mask in training mode, but is the identity in evaluation mode. Args: mode (bool): Indicate if the model should be in training or evaluation mode. Default: ``True``. Returns: The module instance after updating the training mode. """ def eval(self) -> Module: """Set the model to evaluation mode. See :func:`train`. """ def set_dtype( self, dtype: mx.Dtype, predicate: Optional[Callable[[mx.Dtype], bool]] = ... ): # -> None: """Set the dtype of the module's parameters. Args: dtype (Dtype): The new dtype. predicate (typing.Callable, optional): A predicate to select parameters to cast. By default, only parameters of type :attr:`floating` will be updated to avoid casting integer parameters to the new dtype. """ ================================================ FILE: .mlx_typings/mlx/nn/layers/containers.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Callable import mlx.core as mx from base import Module class Sequential(Module): """A layer that calls the passed callables in order. We can pass either modules or plain callables to the Sequential module. If our functions have learnable parameters they should be implemented as ``nn.Module`` instances. Args: modules (tuple of Callables): The modules to call in order """ def __init__(self, *modules: Module | Callable[[mx.array], mx.array]) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mlx/nn/layers/convolution.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Union import mlx.core as mx from base import Module class Conv1d(Module): """Applies a 1-dimensional convolution over the multi-channel input sequence. The channels are expected to be last i.e. the input shape should be ``NLC`` where: * ``N`` is the batch dimension * ``L`` is the sequence length * ``C`` is the number of input channels Args: in_channels (int): The number of input channels out_channels (int): The number of output channels kernel_size (int): The size of the convolution filters stride (int, optional): The stride when applying the filter. Default: ``1``. padding (int, optional): How many positions to 0-pad the input with. Default: ``0``. dilation (int, optional): The dilation of the convolution. groups (int, optional): The number of groups for the convolution. Default: ``1``. bias (bool, optional): If ``True`` add a learnable bias to the output. Default: ``True`` """ weight: mx.array bias: mx.array | None groups: int def __init__( self, in_channels: int, out_channels: int, kernel_size: int, stride: int = ..., padding: int = ..., dilation: int = ..., groups: int = ..., bias: bool = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class Conv2d(Module): """Applies a 2-dimensional convolution over the multi-channel input image. The channels are expected to be last i.e. the input shape should be ``NHWC`` where: * ``N`` is the batch dimension * ``H`` is the input image height * ``W`` is the input image width * ``C`` is the number of input channels Args: in_channels (int): The number of input channels. out_channels (int): The number of output channels. kernel_size (int or tuple): The size of the convolution filters. stride (int or tuple, optional): The size of the stride when applying the filter. Default: ``1``. padding (int or tuple, optional): How many positions to 0-pad the input with. Default: ``0``. dilation (int or tuple, optional): The dilation of the convolution. groups (int, optional): The number of groups for the convolution. Default: ``1``. bias (bool, optional): If ``True`` add a learnable bias to the output. Default: ``True`` """ def __init__( self, in_channels: int, out_channels: int, kernel_size: Union[int, tuple], stride: Union[int, tuple] = ..., padding: Union[int, tuple] = ..., dilation: Union[int, tuple] = ..., groups: int = ..., bias: bool = ..., ) -> None: ... def __call__(self, x) -> mx.array: ... class Conv3d(Module): """Applies a 3-dimensional convolution over the multi-channel input image. The channels are expected to be last i.e. the input shape should be ``NDHWC`` where: * ``N`` is the batch dimension * ``D`` is the input image depth * ``H`` is the input image height * ``W`` is the input image width * ``C`` is the number of input channels Args: in_channels (int): The number of input channels. out_channels (int): The number of output channels. kernel_size (int or tuple): The size of the convolution filters. stride (int or tuple, optional): The size of the stride when applying the filter. Default: ``1``. dilation (int or tuple, optional): The dilation of the convolution. padding (int or tuple, optional): How many positions to 0-pad the input with. Default: ``0``. bias (bool, optional): If ``True`` add a learnable bias to the output. Default: ``True`` """ def __init__( self, in_channels: int, out_channels: int, kernel_size: Union[int, tuple], stride: Union[int, tuple] = ..., padding: Union[int, tuple] = ..., dilation: Union[int, tuple] = ..., bias: bool = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mlx/nn/layers/convolution_transpose.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Union import mlx.core as mx from base import Module class ConvTranspose1d(Module): """Applies a 1-dimensional transposed convolution over the multi-channel input sequence. The channels are expected to be last i.e. the input shape should be ``NLC`` where: * ``N`` is the batch dimension * ``L`` is the sequence length * ``C`` is the number of input channels Args: in_channels (int): The number of input channels out_channels (int): The number of output channels kernel_size (int): The size of the convolution filters stride (int, optional): The stride when applying the filter. Default: ``1``. padding (int, optional): How many positions to 0-pad the input with. Default: ``0``. dilation (int, optional): The dilation of the convolution. output_padding(int, optional): Additional size added to one side of the output shape. Default: ``0``. bias (bool, optional): If ``True`` add a learnable bias to the output. Default: ``True`` """ def __init__( self, in_channels: int, out_channels: int, kernel_size: int, stride: int = ..., padding: int = ..., dilation: int = ..., output_padding: int = ..., bias: bool = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class ConvTranspose2d(Module): """Applies a 2-dimensional transposed convolution over the multi-channel input image. The channels are expected to be last i.e. the input shape should be ``NHWC`` where: * ``N`` is the batch dimension * ``H`` is the input image height * ``W`` is the input image width * ``C`` is the number of input channels Args: in_channels (int): The number of input channels. out_channels (int): The number of output channels. kernel_size (int or tuple): The size of the convolution filters. stride (int or tuple, optional): The size of the stride when applying the filter. Default: ``1``. padding (int or tuple, optional): How many positions to 0-pad the input with. Default: ``0``. dilation (int or tuple, optional): The dilation of the convolution. output_padding(int or tuple, optional): Additional size added to one side of the output shape. Default: ``0``. bias (bool, optional): If ``True`` add a learnable bias to the output. Default: ``True`` """ def __init__( self, in_channels: int, out_channels: int, kernel_size: Union[int, tuple], stride: Union[int, tuple] = ..., padding: Union[int, tuple] = ..., dilation: Union[int, tuple] = ..., output_padding: Union[int, tuple] = ..., bias: bool = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class ConvTranspose3d(Module): """Applies a 3-dimensional transposed convolution over the multi-channel input image. The channels are expected to be last i.e. the input shape should be ``NDHWC`` where: * ``N`` is the batch dimension * ``D`` is the input image depth * ``H`` is the input image height * ``W`` is the input image width * ``C`` is the number of input channels Args: in_channels (int): The number of input channels. out_channels (int): The number of output channels. kernel_size (int or tuple): The size of the convolution filters. stride (int or tuple, optional): The size of the stride when applying the filter. Default: ``1``. padding (int or tuple, optional): How many positions to 0-pad the input with. Default: ``0``. dilation (int or tuple, optional): The dilation of the convolution. output_padding(int or tuple, optional): Additional size added to one side of the output shape. Default: ``0``. bias (bool, optional): If ``True`` add a learnable bias to the output. Default: ``True`` """ def __init__( self, in_channels: int, out_channels: int, kernel_size: Union[int, tuple], stride: Union[int, tuple] = ..., padding: Union[int, tuple] = ..., dilation: Union[int, tuple] = ..., output_padding: Union[int, tuple] = ..., bias: bool = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mlx/nn/layers/distributed.pyi ================================================ """ This type stub file was generated by pyright. """ from functools import lru_cache from typing import Callable, Optional, Union import mlx.core as mx from base import Module from mlx.nn.layers.linear import Linear @lru_cache def sum_gradients( group: mx.distributed.Group, ) -> Callable[..., mx.array]: # -> Callable[..., Any] | Callable[..., array]: ... def shard_inplace( module: Module, sharding: str, *, segments: Union[int, list[int]] = ..., group: Optional[mx.distributed.Group] = ..., ) -> None: """Shard a module in-place by updating its parameter dictionary with the sharded parameter dictionary. The ``sharding`` argument can be any callable that given the path and the weight returns the sharding axis and optionally also the segments that comprise the unsharded weight. For instance if the weight is a fused QKV matrix the segments should be 3. .. note:: The module doesn't change so in order for distributed communication to happen the module needs to natively support it and for it to be enabled. Args: module (Module): The parameters of this module will be sharded in-place. sharding (str or callable): One of "all-to-sharded" and "sharded-to-all" or a callable that returns the sharding axis and segments. segments (int or list): The segments to use if ``sharding`` is a string. Default: ``1``. group (mlx.core.distributed.Group): The distributed group to shard across. If not set, the global group will be used. Default: ``None``. """ def shard_linear( module: Module, sharding: str, *, segments: Union[int, list[int]] = ..., group: Optional[mx.distributed.Group] = ..., ) -> Linear: """Create a new linear layer that has its parameters sharded and also performs distributed communication either in the forward or backward pass. .. note:: Contrary to ``shard_inplace``, the original layer is not changed but a new layer is returned. Args: module (Module): The linear layer to be sharded. sharding (str): One of "all-to-sharded" and "sharded-to-all" that defines the type of sharding to perform. segments (int or list): The segments to use. Default: ``1``. group (mlx.core.distributed.Group): The distributed group to shard across. If not set, the global group will be used. Default: ``None``. """ class AllToShardedLinear(Module): """Each member of the group applies part of the affine transformation such that the result is sharded across the group. The gradients are automatically aggregated from each member of the group. Args: input_dims (int): The dimensionality of the input features output_dims (int): The dimensionality of the output features bias (bool, optional): If set to ``False`` the the layer will not use a bias. Default is ``True``. group (mx.distributed.Group, optional): The sharding will happen across this group. If not set then the global group is used. Default is ``None``. """ def __init__( self, input_dims: int, output_dims: int, bias: bool = ..., group: Optional[mx.distributed.Group] = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... @classmethod def from_linear( cls, linear_layer: Module, *, segments: Union[int, list[int]] = ..., group: Optional[mx.distributed.Group] = ..., ) -> AllToShardedLinear: ... class ShardedToAllLinear(Module): """Each member of the group applies part of the affine transformation and then aggregates the results. All nodes will have the same exact result after this layer. :class:`ShardedToAllLinear` provides a classmethod :meth:`from_linear` to convert linear layers to sharded :obj:`ShardedToAllLinear` layers. Args: input_dims (int): The dimensionality of the input features output_dims (int): The dimensionality of the output features bias (bool, optional): If set to ``False`` the the layer will not use a bias. Default is ``True``. group (mx.distributed.Group, optional): The sharding will happen across this group. If not set then the global group is used. Default is ``None``. """ def __init__( self, input_dims: int, output_dims: int, bias: bool = ..., group: Optional[mx.distributed.Group] = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... @classmethod def from_linear( cls, linear_layer: Module, *, segments: Union[int, list[int]] = ..., group: Optional[mx.distributed.Group] = ..., ) -> ShardedToAllLinear: ... class QuantizedAllToShardedLinear(Module): """Each member of the group applies part of the affine transformation with a quantized matrix such that the result is sharded across the group. It is the quantized equivalent of :class:`AllToShardedLinear`. Similar to :class:`QuantizedLinear` its parameters are frozen and will not be included in any gradient computation. Args: input_dims (int): The dimensionality of the input features. output_dims (int): The dimensionality of the output features. bias (bool, optional): If set to ``False`` then the layer will not use a bias. Default: ``True``. group_size (int, optional): The group size to use for the quantized weight. See :func:`~mlx.core.quantize`. Default: ``64``. bits (int, optional): The bit width to use for the quantized weight. See :func:`~mlx.core.quantize`. Default: ``4``. group (mx.distributed.Group, optional): The sharding will happen across this group. If not set then the global group is used. Default is ``None``. """ def __init__( self, input_dims: int, output_dims: int, bias: bool = ..., group_size: int = ..., bits: int = ..., group: Optional[mx.distributed.Group] = ..., ) -> None: ... def unfreeze(self, *args, **kwargs) -> None: """Wrap unfreeze so that we unfreeze any layers we might contain but our parameters will remain frozen.""" def __call__(self, x: mx.array) -> mx.array: ... @classmethod def from_quantized_linear( cls, quantized_linear_layer: Module, *, segments: Union[int, list[int]] = ..., group: Optional[mx.distributed.Group] = ..., ) -> QuantizedAllToShardedLinear: ... class QuantizedShardedToAllLinear(Module): """Each member of the group applies part of the affine transformation using the quantized matrix and then aggregates the results. All nodes will have the same exact result after this layer. It is the quantized equivalent of :class:`ShardedToAllLinear`. Similar to :class:`QuantizedLinear` its parameters are frozen and will not be included in any gradient computation. Args: input_dims (int): The dimensionality of the input features. output_dims (int): The dimensionality of the output features. bias (bool, optional): If set to ``False`` then the layer will not use a bias. Default: ``True``. group_size (int, optional): The group size to use for the quantized weight. See :func:`~mlx.core.quantize`. Default: ``64``. bits (int, optional): The bit width to use for the quantized weight. See :func:`~mlx.core.quantize`. Default: ``4``. group (mx.distributed.Group, optional): The sharding will happen across this group. If not set then the global group is used. Default is ``None``. """ def __init__( self, input_dims: int, output_dims: int, bias: bool = ..., group_size: int = ..., bits: int = ..., group: Optional[mx.distributed.Group] = ..., ) -> None: ... def unfreeze(self, *args, **kwargs): # -> None: """Wrap unfreeze so that we unfreeze any layers we might contain but our parameters will remain frozen.""" def __call__(self, x: mx.array) -> mx.array: ... @classmethod def from_quantized_linear( cls, quantized_linear_layer: Module, *, segments: Union[int, list[int]] = ..., group: Optional[mx.distributed.Group] = ..., ) -> QuantizedShardedToAllLinear: ... ================================================ FILE: .mlx_typings/mlx/nn/layers/dropout.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from base import Module class Dropout(Module): r"""Randomly zero a portion of the elements during training. The remaining elements are multiplied with :math:`\frac{1}{1-p}` where :math:`p` is the probability of zeroing an element. This is done so the expected value of a given element will remain the same. Args: p (float): The probability to zero an element """ def __init__(self, p: float = ...) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class Dropout2d(Module): r"""Apply 2D channel-wise dropout during training. Randomly zero out entire channels independently with probability :math:`p`. This layer expects the channels to be last, i.e. the input shape should be ``NWHC`` or ``WHC`` where:``N`` is the batch dimension,``H`` is the input image height,``W`` is the input image width, and``C`` is the number of input channels The remaining channels are scaled by :math:`\frac{1}{1-p}` to maintain the expected value of each element. Unlike traditional dropout, which zeros individual entries, this layer zeros entire channels. This is beneficial for early convolution layers where adjacent pixels are correlated. In such case, traditional dropout may not effectively regularize activations. For more details, see [1]. [1]: Thompson, J., Goroshin, R., Jain, A., LeCun, Y. and Bregler C., 2015. Efficient Object Localization Using Convolutional Networks. CVPR 2015. Args: p (float): Probability of zeroing a channel during training. """ def __init__(self, p: float = ...) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class Dropout3d(Module): r"""Apply 3D channel-wise dropout during training. Randomly zero out entire channels independently with probability :math:`p`. This layer expects the channels to be last, i.e., the input shape should be `NDHWC` or `DHWC` where: `N` is the batch dimension, `D` is the depth, `H` is the input image height, `W` is the input image width, and `C` is the number of input channels. The remaining channels are scaled by :math:`\frac{1}{1-p}` to maintain the expected value of each element. Unlike traditional dropout, which zeros individual entries, this layer zeros entire channels. This is often beneficial for convolutional layers processing 3D data, like in medical imaging or video processing. Args: p (float): Probability of zeroing a channel during training. """ def __init__(self, p: float = ...) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mlx/nn/layers/embedding.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from base import Module from .quantized import QuantizedEmbedding class Embedding(Module): """Implements a simple lookup table that maps each input integer to a high-dimensional vector. Typically used to embed discrete tokens for processing by neural networks. Args: num_embeddings (int): How many possible discrete tokens can we embed. Usually called the vocabulary size. dims (int): The dimensionality of the embeddings. """ def __init__(self, num_embeddings: int, dims: int) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... def as_linear(self, x: mx.array) -> mx.array: """ Call the embedding layer as a linear layer. Use this for example when input embedding and output projection weights are tied. """ def to_quantized( self, group_size: int = ..., bits: int = ..., mode: str = ... ) -> QuantizedEmbedding: """Return a :obj:`QuantizedEmbedding` layer that approximates this embedding layer.""" ================================================ FILE: .mlx_typings/mlx/nn/layers/linear.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Any import mlx.core as mx from base import Module from .quantized import QuantizedLinear class Identity(Module): r"""A placeholder identity operator that is argument-insensitive. Args: args: any argument (unused) kwargs: any keyword argument (unused) """ def __init__(self, *args: Any, **kwargs: Any) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class Linear(Module): r"""Applies an affine transformation to the input. Concretely: .. math:: y = x W^\top + b where: where :math:`W` has shape ``[output_dims, input_dims]`` and :math:`b` has shape ``[output_dims]``. The values are initialized from the uniform distribution :math:`\mathcal{U}(-{k}, {k})`, where :math:`k = \frac{1}{\sqrt{D_i}}` and :math:`D_i` is equal to ``input_dims``. Args: input_dims (int): The dimensionality of the input features output_dims (int): The dimensionality of the output features bias (bool, optional): If set to ``False`` then the layer will not use a bias. Default is ``True``. """ weight: mx.array bias: mx.array | None def __init__(self, input_dims: int, output_dims: int, bias: bool = ...) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... def to_quantized( self, group_size: int = ..., bits: int = ..., mode: str = ... ) -> QuantizedLinear: """Return a :obj:`QuantizedLinear` layer that approximates this layer.""" class Bilinear(Module): r"""Applies a bilinear transformation to the inputs. Concretely: .. math:: y_i = x_1^\top W_i x_2 + b_i where: :math:`W` has shape ``[output_dims, input1_dims, input2_dims]``, :math:`b` has shape ``[output_dims ]``, and :math:`i` indexes the output dimension. The values are initialized from the uniform distribution :math:`\mathcal{U}(-{k}, {k})`, where :math:`k = \frac{1}{\sqrt{D_1}}` and :math:`D_1` is ``input1_dims``. Args: input1_dims (int): The dimensionality of the input1 features input2_dims (int): The dimensionality of the input2 features output_dims (int): The dimensionality of the output features bias (bool, optional): If set to ``False`` then the layer will not use a bias. Default is ``True``. """ def __init__( self, input1_dims: int, input2_dims: int, output_dims: int, bias: bool = ... ) -> None: ... def __call__(self, x1: mx.array, x2: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mlx/nn/layers/normalization.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.core as mx from base import Module class InstanceNorm(Module): r"""Applies instance normalization [1] on the inputs. Computes .. math:: y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta, where :math:`\gamma` and :math:`\beta` are learned per feature dimension parameters initialized at 1 and 0 respectively. Both are of size :attr:`dims`, if :attr:`affine` is ``True``. Args: dims (int): The number of features of the input. eps (float): A value added to the denominator for numerical stability. Default: ``1e-5``. affine (bool): Default: ``False``. Shape: - Input: :math:`(..., C)` where :math:`C` is equal to :attr:`dims`. - Output: Same shape as the input. Examples: >>> import mlx.core as mx >>> import mlx.nn as nn >>> x = mx.random.normal((8, 4, 4, 16)) >>> inorm = nn.InstanceNorm(dims=16) >>> output = inorm(x) References: [1]: https://arxiv.org/abs/1607.08022 """ def __init__(self, dims: int, eps: float = ..., affine: bool = ...) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class LayerNorm(Module): r"""Applies layer normalization [1] on the inputs. Computes .. math:: y = \frac{x - E[x]}{\sqrt{Var[x]} + \epsilon} \gamma + \beta, where :math:`\gamma` and :math:`\beta` are learned per feature dimension parameters initialized at 1 and 0 respectively. [1]: https://arxiv.org/abs/1607.06450 Args: dims (int): The feature dimension of the input to normalize over eps (float): A small additive constant for numerical stability affine (bool): If True learn an affine transform to apply after the normalization bias (bool): If True include a translation to the affine transformation. If set to False the transformation is not really affine just scaling. """ def __init__( self, dims: int, eps: float = ..., affine: bool = ..., bias: bool = ... ) -> None: ... def __call__(self, x) -> mx.array: ... class RMSNorm(Module): r"""Applies Root Mean Square normalization [1] to the inputs. Computes .. math:: y = \frac{x}{\sqrt{E[x^2] + \epsilon}} \gamma where :math:`\gamma` is a learned per feature dimension parameter initialized at 1. Note the accumulation for the mean is done in 32-bit precision. [1]: https://arxiv.org/abs/1910.07467 Args: dims (int): The feature dimension of the input to normalize over eps (float): A small additive constant for numerical stability """ weight: mx.array def __init__(self, dims: int, eps: float = ...) -> None: ... def __call__(self, x) -> mx.array: ... class GroupNorm(Module): r"""Applies Group Normalization [1] to the inputs. Computes the same normalization as layer norm, namely .. math:: y = \frac{x - E[x]}{\sqrt{Var[x]} + \epsilon} \gamma + \beta, where :math:`\gamma` and :math:`\beta` are learned per feature dimension parameters initialized at 1 and 0 respectively. However, the mean and variance are computed over the spatial dimensions and each group of features. In particular, the input is split into num_groups across the feature dimension. The feature dimension is assumed to be the last dimension and the dimensions that precede it (except the first) are considered the spatial dimensions. [1]: https://arxiv.org/abs/1803.08494 Args: num_groups (int): Number of groups to separate the features into dims (int): The feature dimensions of the input to normalize over eps (float): A small additive constant for numerical stability affine (bool): If True learn an affine transform to apply after the normalization. pytorch_compatible (bool): If True perform the group normalization in the same order/grouping as PyTorch. """ def __init__( self, num_groups: int, dims: int, eps: float = ..., affine: bool = ..., pytorch_compatible: bool = ..., ) -> None: ... def __call__(self, x) -> mx.array: ... class BatchNorm(Module): r"""Applies Batch Normalization over a 2D or 3D input. Computes .. math:: y = \frac{x - E[x]}{\sqrt{Var[x]} + \epsilon} \gamma + \beta, where :math:`\gamma` and :math:`\beta` are learned per feature dimension parameters initialized at 1 and 0 respectively. The input shape is specified as ``NC`` or ``NLC``, where ``N`` is the batch, ``C`` is the number of features or channels, and ``L`` is the sequence length. The output has the same shape as the input. For four-dimensional arrays, the shape is ``NHWC``, where ``H`` and ``W`` are the height and width respectively. For more information on Batch Normalization, see the original paper `Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift `_. Args: num_features (int): The feature dimension to normalize over. eps (float, optional): A small additive constant for numerical stability. Default: ``1e-5``. momentum (float, optional): The momentum for updating the running mean and variance. Default: ``0.1``. affine (bool, optional): If ``True``, apply a learned affine transformation after the normalization. Default: ``True``. track_running_stats (bool, optional): If ``True``, track the running mean and variance. Default: ``True``. Examples: >>> import mlx.core as mx >>> import mlx.nn as nn >>> x = mx.random.normal((5, 4)) >>> bn = nn.BatchNorm(num_features=4, affine=True) >>> output = bn(x) """ def __init__( self, num_features: int, eps: float = ..., momentum: float = ..., affine: bool = ..., track_running_stats: bool = ..., ) -> None: ... def unfreeze(self, *args, **kwargs): # -> None: """Wrap unfreeze to make sure that running_mean and var are always frozen parameters.""" def __call__(self, x: mx.array) -> mx.array: """ Forward pass of BatchNorm. Args: x (array): Input tensor. Returns: array: Normalized output tensor. """ ================================================ FILE: .mlx_typings/mlx/nn/layers/pooling.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Optional, Tuple, Union import mlx.core as mx from base import Module class _Pool(Module): def __init__( self, pooling_function, kernel_size, stride, padding, padding_value ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class _Pool1d(_Pool): def __init__( self, pooling_function, padding_value, kernel_size: Union[int, Tuple[int]], stride: Optional[Union[int, Tuple[int]]] = ..., padding: Union[int, Tuple[int]] = ..., ) -> None: ... class _Pool2d(_Pool): def __init__( self, pooling_function, padding_value, kernel_size: Union[int, Tuple[int, int]], stride: Optional[Union[int, Tuple[int, int]]] = ..., padding: Optional[Union[int, Tuple[int, int]]] = ..., ) -> None: ... class _Pool3d(_Pool): def __init__( self, pooling_function, padding_value, kernel_size: Union[int, Tuple[int, int, int]], stride: Optional[Union[int, Tuple[int, int, int]]] = ..., padding: Optional[Union[int, Tuple[int, int, int]]] = ..., ) -> None: ... class MaxPool1d(_Pool1d): r"""Applies 1-dimensional max pooling. Spatially downsamples the input by taking the maximum of a sliding window of size ``kernel_size`` and sliding stride ``stride``. Args: kernel_size (int or tuple(int)): The size of the pooling window kernel. stride (int or tuple(int), optional): The stride of the pooling window. Default: ``kernel_size``. padding (int or tuple(int), optional): How much negative infinity padding to apply to the input. The padding amount is applied to both sides of the spatial axis. Default: ``0``. Examples: >>> import mlx.core as mx >>> import layers as nn >>> x = mx.random.normal(shape=(4, 16, 5)) >>> pool = nn.MaxPool1d(kernel_size=2, stride=2) >>> pool(x) """ def __init__( self, kernel_size: Union[int, Tuple[int]], stride: Optional[Union[int, Tuple[int]]] = ..., padding: Union[int, Tuple[int]] = ..., ) -> None: ... class AvgPool1d(_Pool1d): r"""Applies 1-dimensional average pooling. Spatially downsamples the input by taking the average of a sliding window of size ``kernel_size`` and sliding stride ``stride``. Args: kernel_size (int or tuple(int)): The size of the pooling window kernel. stride (int or tuple(int), optional): The stride of the pooling window. Default: ``kernel_size``. padding (int or tuple(int), optional): How much zero padding to apply to the input. The padding amount is applied to both sides of the spatial axis. Default: ``0``. Examples: >>> import mlx.core as mx >>> import layers as nn >>> x = mx.random.normal(shape=(4, 16, 5)) >>> pool = nn.AvgPool1d(kernel_size=2, stride=2) >>> pool(x) """ def __init__( self, kernel_size: Union[int, Tuple[int]], stride: Optional[Union[int, Tuple[int]]] = ..., padding: Union[int, Tuple[int]] = ..., ) -> None: ... class MaxPool2d(_Pool2d): r"""Applies 2-dimensional max pooling. Spatially downsamples the input by taking the maximum of a sliding window of size ``kernel_size`` and sliding stride ``stride``. The parameters ``kernel_size``, ``stride``, and ``padding`` can either be: * a single ``int`` -- in which case the same value is used for both the height and width axis. * a ``tuple`` of two ``int`` s -- in which case, the first ``int`` is used for the height axis, the second ``int`` for the width axis. Args: kernel_size (int or tuple(int, int)): The size of the pooling window. stride (int or tuple(int, int), optional): The stride of the pooling window. Default: ``kernel_size``. padding (int or tuple(int, int), optional): How much negative infinity padding to apply to the input. The padding is applied on both sides of the height and width axis. Default: ``0``. Examples: >>> import mlx.core as mx >>> import layers as nn >>> x = mx.random.normal(shape=(8, 32, 32, 4)) >>> pool = nn.MaxPool2d(kernel_size=2, stride=2) >>> pool(x) """ def __init__( self, kernel_size: Union[int, Tuple[int, int]], stride: Optional[Union[int, Tuple[int, int]]] = ..., padding: Optional[Union[int, Tuple[int, int]]] = ..., ) -> None: ... class AvgPool2d(_Pool2d): r"""Applies 2-dimensional average pooling. Spatially downsamples the input by taking the average of a sliding window of size ``kernel_size`` and sliding stride ``stride``. The parameters ``kernel_size``, ``stride``, and ``padding`` can either be: * a single ``int`` -- in which case the same value is used for both the height and width axis. * a ``tuple`` of two ``int`` s -- in which case, the first ``int`` is used for the height axis, the second ``int`` for the width axis. Args: kernel_size (int or tuple(int, int)): The size of the pooling window. stride (int or tuple(int, int), optional): The stride of the pooling window. Default: ``kernel_size``. padding (int or tuple(int, int), optional): How much zero padding to apply to the input. The padding is applied on both sides of the height and width axis. Default: ``0``. Examples: >>> import mlx.core as mx >>> import layers as nn >>> x = mx.random.normal(shape=(8, 32, 32, 4)) >>> pool = nn.AvgPool2d(kernel_size=2, stride=2) >>> pool(x) """ def __init__( self, kernel_size: Union[int, Tuple[int, int]], stride: Optional[Union[int, Tuple[int, int]]] = ..., padding: Optional[Union[int, Tuple[int, int]]] = ..., ) -> None: ... class MaxPool3d(_Pool3d): r"""Applies 3-dimensional max pooling. Spatially downsamples the input by taking the maximum of a sliding window of size ``kernel_size`` and sliding stride ``stride``. The parameters ``kernel_size``, ``stride``, and ``padding`` can either be: * a single ``int`` -- in which case the same value is used for the depth, height, and width axis. * a ``tuple`` of three ``int`` s -- in which case, the first ``int`` is used for the depth axis, the second ``int`` for the height axis, and the third ``int`` for the width axis. Args: kernel_size (int or tuple(int, int, int)): The size of the pooling window. stride (int or tuple(int, int, int), optional): The stride of the pooling window. Default: ``kernel_size``. padding (int or tuple(int, int, int), optional): How much negative infinity padding to apply to the input. The padding is applied on both sides of the depth, height and width axis. Default: ``0``. Examples: >>> import mlx.core as mx >>> import layers as nn >>> x = mx.random.normal(shape=(8, 16, 32, 32, 4)) >>> pool = nn.MaxPool3d(kernel_size=2, stride=2) >>> pool(x) """ def __init__( self, kernel_size: Union[int, Tuple[int, int, int]], stride: Optional[Union[int, Tuple[int, int, int]]] = ..., padding: Optional[Union[int, Tuple[int, int, int]]] = ..., ) -> None: ... class AvgPool3d(_Pool3d): r"""Applies 3-dimensional average pooling. Spatially downsamples the input by taking the average of a sliding window of size ``kernel_size`` and sliding stride ``stride``. The parameters ``kernel_size``, ``stride``, and ``padding`` can either be: * a single ``int`` -- in which case the same value is used for the depth, height, and width axis. * a ``tuple`` of three ``int`` s -- in which case, the first ``int`` is used for the depth axis, the second ``int`` for the height axis, and the third ``int`` for the width axis. Args: kernel_size (int or tuple(int, int, int)): The size of the pooling window. stride (int or tuple(int, int, int), optional): The stride of the pooling window. Default: ``kernel_size``. padding (int or tuple(int, int, int), optional): How much zero padding to apply to the input. The padding is applied on both sides of the depth, height and width axis. Default: ``0``. Examples: >>> import mlx.core as mx >>> import layers as nn >>> x = mx.random.normal(shape=(8, 16, 32, 32, 4)) >>> pool = nn.AvgPool3d(kernel_size=2, stride=2) >>> pool(x) """ def __init__( self, kernel_size: Union[int, Tuple[int, int, int]], stride: Optional[Union[int, Tuple[int, int, int]]] = ..., padding: Optional[Union[int, Tuple[int, int, int]]] = ..., ) -> None: ... ================================================ FILE: .mlx_typings/mlx/nn/layers/positional_encoding.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Optional import mlx.core as mx from base import Module class RoPE(Module): """Implements the rotary positional encoding. The traditional implementation rotates consecutive pairs of elements in the feature dimension while the default implementation rotates pairs with stride half the feature dimensions for efficiency. For more details see `RoFormer: Enhanced Transformer with Rotary Position Embedding `_. Args: dims (int): The feature dimensions to be rotated. If the input feature is larger than dims then the rest is left unchanged. traditional (bool, optional): If set to ``True`` choose the traditional implementation which is slightly less efficient. Default: ``False``. base (float, optional): The base used to compute angular frequency for each dimension in the positional encodings. Default: ``10000``. scale (float, optional): The scale used to scale the positions. Default: ``1.0``. """ def __init__( self, dims: int, traditional: bool = ..., base: float = ..., scale: float = ... ) -> None: ... def __call__(self, x, offset: int = ...) -> mx.array: ... class SinusoidalPositionalEncoding(Module): r"""Implements sinusoidal positional encoding. For more details see the paper `Attention Is All You Need `_. Args: dims (int): The dimensionality of the resulting positional embeddings. min_freq (float, optional): The minimum frequency expected. Default: ``0.0001``. max_freq (float, optional): The maximum frequency expected. Default: ``1``. scale (float, optional): A multiplicative scale for the embeddings. Default: ``sqrt(2/dims)``. cos_first (bool, optional): If ``True`` embed using ``[cos(x); sin(x)]`` instead of the reverse. Default: ``False``. full_turns (bool, optional): If ``True`` multiply the frequencies with :math:`2\pi`. Default: ``False``. """ def __init__( self, dims: int, min_freq: float = ..., max_freq: float = ..., scale: Optional[float] = ..., cos_first: bool = ..., full_turns: bool = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class ALiBi(Module): _alibi_mask_key = ... _alibi_mask = ... @classmethod def create_alibi_matrix( cls, q_sequence_length: int, k_sequence_length: int, num_heads: int, offset: int, dtype=..., ) -> mx.array | None: ... @staticmethod def create_alibi_slope(num_heads: int) -> mx.array: ... def __call__( self, attention_scores: mx.array, offset=..., mask=... ) -> mx.array: ... ================================================ FILE: .mlx_typings/mlx/nn/layers/quantized.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Callable, Optional, Union import mlx.core as mx from base import Module def quantize( model: Module, group_size: int = ..., bits: int = ..., *, mode: str = ..., class_predicate: Optional[Callable[[str, Module], Union[bool, dict]]] = ..., ): # -> None: """Quantize the sub-modules of a module according to a predicate. By default all layers that define a ``to_quantized(group_size, bits)`` method will be quantized. Both :obj:`Linear` and :obj:`Embedding` layers will be quantized. Note also, the module is updated in-place. Args: model (Module): The model whose leaf modules may be quantized. group_size (int): The quantization group size (see :func:`mlx.core.quantize`). Default: ``64``. bits (int): The number of bits per parameter (see :func:`mlx.core.quantize`). Default: ``4``. mode (str): The quantization method to use (see :func:`mlx.core.quantize`). Default: ``"affine"``. class_predicate (Optional[Callable]): A callable which receives the :obj:`Module` path and :obj:`Module` itself and returns ``True`` or a dict of params for `to_quantized` if it should be quantized and ``False`` otherwise. If ``None``, then all layers that define a ``to_quantized(group_size, bits)`` method are quantized. Default: ``None``. """ class QuantizedEmbedding(Module): """The same as :obj:`Embedding` but with a quantized weight matrix. :obj:`QuantizedEmbedding` also provides a :meth:`from_embedding` classmethod to convert embedding layers to :obj:`QuantizedEmbedding` layers. Args: num_embeddings (int): How many possible discrete tokens can we embed. Usually called the vocabulary size. dims (int): The dimensionality of the embeddings. group_size (int, optional): The group size to use for the quantized weight. See :func:`~mlx.core.quantize`. Default: ``64``. bits (int, optional): The bit width to use for the quantized weight. See :func:`~mlx.core.quantize`. Default: ``4``. mode (str): The quantization method to use (see :func:`mlx.core.quantize`). Default: ``"affine"``. """ def __init__( self, num_embeddings: int, dims: int, group_size: int = ..., bits: int = ..., mode: str = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... def as_linear(self, x: mx.array) -> mx.array: """ Call the quantized embedding layer as a quantized linear layer. Use this for example when input embedding and output projection weights are tied. """ @classmethod def from_embedding( cls, embedding_layer: Module, group_size: int = ..., bits: int = ..., mode: str = ..., ) -> QuantizedEmbedding: """Create a :obj:`QuantizedEmbedding` layer from an :obj:`Embedding` layer.""" class QuantizedLinear(Module): """Applies an affine transformation to the input using a quantized weight matrix. It is the quantized equivalent of :class:`Linear`. For now its parameters are frozen and will not be included in any gradient computation but this will probably change in the future. :obj:`QuantizedLinear` also provides a classmethod :meth:`from_linear` to convert linear layers to :obj:`QuantizedLinear` layers. Args: input_dims (int): The dimensionality of the input features. output_dims (int): The dimensionality of the output features. bias (bool, optional): If set to ``False`` then the layer will not use a bias. Default: ``True``. group_size (int, optional): The group size to use for the quantized weight. See :func:`~mlx.core.quantize`. Default: ``64``. bits (int, optional): The bit width to use for the quantized weight. See :func:`~mlx.core.quantize`. Default: ``4``. mode (str): The quantization method to use (see :func:`mlx.core.quantize`). Default: ``"affine"``. """ def __init__( self, input_dims: int, output_dims: int, bias: bool = ..., group_size: int = ..., bits: int = ..., mode: str = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... @classmethod def from_linear( cls, linear_layer: Module, group_size: int = ..., bits: int = ..., mode: str = ..., ) -> QuantizedLinear: """Create a :obj:`QuantizedLinear` layer from a :obj:`Linear` layer.""" ================================================ FILE: .mlx_typings/mlx/nn/layers/recurrent.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Callable, Optional import mlx.core as mx from base import Module class RNN(Module): r"""An Elman recurrent layer. The input is a sequence of shape ``NLD`` or ``LD`` where: * ``N`` is the optional batch dimension * ``L`` is the sequence length * ``D`` is the input's feature dimension Concretely, for each element along the sequence length axis, this layer applies the function: .. math:: h_{t + 1} = \text{tanh} (W_{ih}x_t + W_{hh}h_t + b) The hidden state :math:`h` has shape ``NH`` or ``H``, depending on whether the input is batched or not. Returns the hidden state at each time step, of shape ``NLH`` or ``LH``. Args: input_size (int): Dimension of the input, ``D``. hidden_size (int): Dimension of the hidden state, ``H``. bias (bool, optional): Whether to use a bias. Default: ``True``. nonlinearity (callable, optional): Non-linearity to use. If ``None``, then func:`tanh` is used. Default: ``None``. """ def __init__( self, input_size: int, hidden_size: int, bias: bool = ..., nonlinearity: Optional[Callable] = ..., ) -> None: ... def __call__(self, x: mx.array, hidden=...) -> mx.array: ... class GRU(Module): r"""A gated recurrent unit (GRU) RNN layer. The input has shape ``NLD`` or ``LD`` where: * ``N`` is the optional batch dimension * ``L`` is the sequence length * ``D`` is the input's feature dimension Concretely, for each element of the sequence, this layer computes: .. math:: \begin{aligned} r_t &= \sigma (W_{xr}x_t + W_{hr}h_t + b_{r}) \\ z_t &= \sigma (W_{xz}x_t + W_{hz}h_t + b_{z}) \\ n_t &= \text{tanh}(W_{xn}x_t + b_{n} + r_t \odot (W_{hn}h_t + b_{hn})) \\ h_{t + 1} &= (1 - z_t) \odot n_t + z_t \odot h_t \end{aligned} The hidden state :math:`h` has shape ``NH`` or ``H`` depending on whether the input is batched or not. Returns the hidden state at each time step of shape ``NLH`` or ``LH``. Args: input_size (int): Dimension of the input, ``D``. hidden_size (int): Dimension of the hidden state, ``H``. bias (bool): Whether to use biases or not. Default: ``True``. """ def __init__(self, input_size: int, hidden_size: int, bias: bool = ...) -> None: ... def __call__(self, x: mx.array, hidden=...) -> mx.array: ... class LSTM(Module): r"""An LSTM recurrent layer. The input has shape ``NLD`` or ``LD`` where: * ``N`` is the optional batch dimension * ``L`` is the sequence length * ``D`` is the input's feature dimension Concretely, for each element of the sequence, this layer computes: .. math:: \begin{aligned} i_t &= \sigma (W_{xi}x_t + W_{hi}h_t + b_{i}) \\ f_t &= \sigma (W_{xf}x_t + W_{hf}h_t + b_{f}) \\ g_t &= \text{tanh} (W_{xg}x_t + W_{hg}h_t + b_{g}) \\ o_t &= \sigma (W_{xo}x_t + W_{ho}h_t + b_{o}) \\ c_{t + 1} &= f_t \odot c_t + i_t \odot g_t \\ h_{t + 1} &= o_t \text{tanh}(c_{t + 1}) \end{aligned} The hidden state :math:`h` and cell state :math:`c` have shape ``NH`` or ``H``, depending on whether the input is batched or not. The layer returns two arrays, the hidden state and the cell state at each time step, both of shape ``NLH`` or ``LH``. Args: input_size (int): Dimension of the input, ``D``. hidden_size (int): Dimension of the hidden state, ``H``. bias (bool): Whether to use biases or not. Default: ``True``. """ def __init__(self, input_size: int, hidden_size: int, bias: bool = ...) -> None: ... def __call__( self, x: mx.array, hidden=..., cell=... ) -> tuple[mx.array, mx.array]: ... ================================================ FILE: .mlx_typings/mlx/nn/layers/transformer.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Any, Callable, Optional import mlx.core as mx from base import Module class MultiHeadAttention(Module): """Implements the scaled dot product attention with multiple heads. Given inputs for queries, keys and values the ``MultiHeadAttention`` produces new values by aggregating information from the input values according to the similarities of the input queries and keys. All inputs as well as the output are linearly projected without biases by default. ``MultiHeadAttention`` also takes an optional additive attention mask that should be broadcastable with ``(batch, num_heads, # queries, # keys)``. The mask should have ``-inf`` or very large negative numbers at the positions that should *not* be attended to. Args: dims (int): The model dimensions. This is also the default value for the queries, keys, values, and the output. num_heads (int): The number of attention heads to use. query_input_dims (int, optional): The input dimensions of the queries. Default: ``dims``. key_input_dims (int, optional): The input dimensions of the keys. Default: ``dims``. value_input_dims (int, optional): The input dimensions of the values. Default: ``key_input_dims``. value_dims (int, optional): The dimensions of the values after the projection. Default: ``dims``. value_output_dims (int, optional): The dimensions the new values will be projected to. Default: ``dims``. bias (bool, optional): Whether or not to use a bias in the projections. Default: ``False``. """ def __init__( self, dims: int, num_heads: int, query_input_dims: Optional[int] = ..., key_input_dims: Optional[int] = ..., value_input_dims: Optional[int] = ..., value_dims: Optional[int] = ..., value_output_dims: Optional[int] = ..., bias: bool = ..., ) -> None: ... def __call__( self, queries: mx.array, keys: mx.array, values: mx.array, mask: mx.array = ... ) -> mx.array: ... @staticmethod def create_additive_causal_mask(N: int, dtype: mx.Dtype = ...) -> mx.array: ... class TransformerEncoderLayer(Module): def __init__( self, dims: int, num_heads: int, mlp_dims: Optional[int] = ..., dropout: float = ..., activation: Callable[[Any], Any] = ..., norm_first: bool = ..., ) -> None: ... def __call__(self, x: mx.array, mask: mx.array) -> mx.array: ... class TransformerEncoder(Module): def __init__( self, num_layers: int, dims: int, num_heads: int, mlp_dims: Optional[int] = ..., dropout: float = ..., activation=..., norm_first: bool = ..., checkpoint: bool = ..., ) -> None: ... def __call__(self, x: mx.array, mask: mx.array) -> mx.array: ... class TransformerDecoderLayer(Module): def __init__( self, dims: int, num_heads: int, mlp_dims: Optional[int] = ..., dropout: float = ..., activation: Callable[[Any], Any] = ..., norm_first: bool = ..., ) -> None: ... def __call__(self, x: mx.array, memory, x_mask, memory_mask) -> mx.array: ... class TransformerDecoder(Module): def __init__( self, num_layers: int, dims: int, num_heads: int, mlp_dims: Optional[int] = ..., dropout: float = ..., activation=..., norm_first: bool = ..., checkpoint: bool = ..., ) -> None: ... def __call__(self, x: mx.array, memory, x_mask, memory_mask) -> mx.array: ... class Transformer(Module): """ Implements a standard Transformer model. The implementation is based on `Attention Is All You Need `_. The Transformer model contains an encoder and a decoder. The encoder processes the input sequence and the decoder generates the output sequence. The interaction between encoder and decoder happens through the attention mechanism. Args: dims (int, optional): The number of expected features in the encoder/decoder inputs. Default: ``512``. num_heads (int, optional): The number of attention heads. Default: ``8``. num_encoder_layers (int, optional): The number of encoder layers in the Transformer encoder. Default: ``6``. num_decoder_layers (int, optional): The number of decoder layers in the Transformer decoder. Default: ``6``. mlp_dims (int, optional): The hidden dimension of the MLP block in each Transformer layer. Defaults to ``4*dims`` if not provided. Default: ``None``. dropout (float, optional): The dropout value for the Transformer encoder and decoder. Dropout is used after each attention layer and the activation in the MLP layer. Default: ``0.0``. activation (function, optional): the activation function for the MLP hidden layer. Default: :func:`relu`. custom_encoder (nn.Module, optional): A custom encoder to replace the standard Transformer encoder. Default: ``None``. custom_decoder (nn.Module, optional): A custom decoder to replace the standard Transformer decoder. Default: ``None``. norm_first (bool, optional): if ``True``, encoder and decoder layers will perform layer normalization before attention and MLP operations, otherwise after. Default: ``True``. checkpoint (bool, optional): if ``True`` perform gradient checkpointing to reduce the memory usage at the expense of more computation. Default: ``False``. """ def __init__( self, dims: int = ..., num_heads: int = ..., num_encoder_layers: int = ..., num_decoder_layers: int = ..., mlp_dims: Optional[int] = ..., dropout: float = ..., activation: Callable[[Any], Any] = ..., custom_encoder: Optional[Any] = ..., custom_decoder: Optional[Any] = ..., norm_first: bool = ..., checkpoint: bool = ..., ) -> None: ... def __call__( self, src, tgt, src_mask, tgt_mask, memory_mask ) -> mx.array: # -> array | Any: ... ================================================ FILE: .mlx_typings/mlx/nn/layers/upsample.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Literal, Tuple, Union import mlx.core as mx from base import Module def upsample_nearest(x: mx.array, scale_factor: Tuple) -> mx.array: ... def upsample_linear( x: mx.array, scale_factor: Tuple, align_corners: bool = ... ): # -> int: ... def upsample_cubic( x: mx.array, scale_factor: Tuple, align_corners: bool = ... ): # -> int: ... class Upsample(Module): r"""Upsample the input signal spatially. The spatial dimensions are by convention dimensions ``1`` to ``x.ndim - 2``. The first is the batch dimension and the last is the feature dimension. For example, an audio signal would be 3D with 1 spatial dimension, an image 4D with 2 and so on and so forth. There are three upsampling algorithms implemented nearest neighbor upsampling, linear interpolation, and cubic interpolation. All can be applied to any number of spatial dimensions. The linear interpolation will be bilinear, trilinear etc when applied to more than one spatial dimension. And cubic interpolation will be bicubic when there are 2 spatial dimensions. .. note:: When using one of the linear or cubic interpolation modes the ``align_corners`` argument changes how the corners are treated in the input image. If ``align_corners=True`` then the top and left edge of the input and output will be matching as will the bottom right edge. Parameters: scale_factor (float or tuple): The multiplier for the spatial size. If a ``float`` is provided, it is the multiplier for all spatial dimensions. Otherwise, the number of scale factors provided must match the number of spatial dimensions. mode (str, optional): The upsampling algorithm, either ``"nearest"``, ``"linear"`` or ``"cubic"``. Default: ``"nearest"``. align_corners (bool, optional): Changes the way the corners are treated during ``"linear"`` and ``"cubic"`` upsampling. See the note above and the examples below for more details. Default: ``False``. Examples: >>> import mlx.core as mx >>> import mlx.nn as nn >>> x = mx.arange(1, 5).reshape((1, 2, 2, 1)) >>> x array([[[[1], [2]], [[3], [4]]]], dtype=int32) >>> n = nn.Upsample(scale_factor=2, mode='nearest') >>> n(x).squeeze() array([[1, 1, 2, 2], [1, 1, 2, 2], [3, 3, 4, 4], [3, 3, 4, 4]], dtype=int32) >>> b = nn.Upsample(scale_factor=2, mode='linear') >>> b(x).squeeze() array([[1, 1.25, 1.75, 2], [1.5, 1.75, 2.25, 2.5], [2.5, 2.75, 3.25, 3.5], [3, 3.25, 3.75, 4]], dtype=float32) >>> b = nn.Upsample(scale_factor=2, mode='linear', align_corners=True) >>> b(x).squeeze() array([[1, 1.33333, 1.66667, 2], [1.66667, 2, 2.33333, 2.66667], [2.33333, 2.66667, 3, 3.33333], [3, 3.33333, 3.66667, 4]], dtype=float32) """ def __init__( self, scale_factor: Union[float, Tuple], mode: Literal["nearest", "linear", "cubic"] = ..., align_corners: bool = ..., ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... ================================================ FILE: .mlx_typings/mlx/nn/losses.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Literal, Optional import mlx.core as mx Reduction = Literal["none", "mean", "sum"] def cross_entropy( logits: mx.array, targets: mx.array, weights: Optional[mx.array] = ..., axis: int = ..., label_smoothing: float = ..., reduction: Reduction = ..., ) -> mx.array: """ Computes the cross entropy loss. Args: logits (array): The unnormalized logits. targets (array): The ground truth values. These can be class indices or probabilities for each class. If the ``targets`` are class indices, then ``targets`` shape should match the ``logits`` shape with the ``axis`` dimension removed. If the ``targets`` are probabilities (or one-hot encoded), then the ``targets`` shape should be the same as the ``logits`` shape. weights (array, optional): Optional weights for each target. Default: ``None``. axis (int, optional): The axis over which to compute softmax. Default: ``-1``. label_smoothing (float, optional): Label smoothing factor. Default: ``0``. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: array: The computed cross entropy loss. Examples: >>> import mlx.core as mx >>> import mlx.nn as nn >>> >>> # Class indices as targets >>> logits = mx.array([[2.0, -1.0], [-1.0, 2.0]]) >>> targets = mx.array([0, 1]) >>> nn.losses.cross_entropy(logits, targets) array([0.0485873, 0.0485873], dtype=float32) >>> >>> # Probabilities (or one-hot vectors) as targets >>> logits = mx.array([[2.0, -1.0], [-1.0, 2.0]]) >>> targets = mx.array([[0.9, 0.1], [0.1, 0.9]]) >>> nn.losses.cross_entropy(logits, targets) array([0.348587, 0.348587], dtype=float32) """ def binary_cross_entropy( inputs: mx.array, targets: mx.array, weights: Optional[mx.array] = ..., with_logits: bool = ..., reduction: Reduction = ..., ) -> mx.array: """ Computes the binary cross entropy loss. By default, this function takes the pre-sigmoid logits, which results in a faster and more precise loss. For improved numerical stability when ``with_logits=False``, the loss calculation clips the input probabilities (in log-space) to a minimum value of ``-100``. Args: inputs (array): The predicted values. If ``with_logits`` is ``True``, then ``inputs`` are unnormalized logits. Otherwise, ``inputs`` are probabilities. targets (array): The binary target values in {0, 1}. with_logits (bool, optional): Whether ``inputs`` are logits. Default: ``True``. weights (array, optional): Optional weights for each target. Default: ``None``. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``. Returns: array: The computed binary cross entropy loss. Examples: >>> import mlx.core as mx >>> import mlx.nn as nn >>> logits = mx.array([0.105361, 0.223144, 1.20397, 0.916291]) >>> targets = mx.array([0, 0, 1, 1]) >>> loss = nn.losses.binary_cross_entropy(logits, targets, reduction="mean") >>> loss array(0.539245, dtype=float32) >>> probs = mx.array([0.1, 0.1, 0.4, 0.4]) >>> targets = mx.array([0, 0, 1, 1]) >>> loss = nn.losses.binary_cross_entropy(probs, targets, with_logits=False, reduction="mean") >>> loss array(0.510826, dtype=float32) """ def l1_loss( predictions: mx.array, targets: mx.array, reduction: Reduction = ... ) -> mx.array: """ Computes the L1 loss. Args: predictions (array): The predicted values. targets (array): The target values. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``. Returns: array: The computed L1 loss. """ def mse_loss( predictions: mx.array, targets: mx.array, reduction: Reduction = ... ) -> mx.array: """ Computes the mean squared error loss. Args: predictions (array): The predicted values. targets (array): The target values. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``. Returns: array: The computed mean squared error loss. """ def nll_loss( inputs: mx.array, targets: mx.array, axis: int = ..., reduction: Reduction = ... ) -> mx.array: """ Computes the negative log likelihood loss. Args: inputs (array): The predicted distribution in log space. targets (array): The target values. axis (int, optional): The distribution axis. Default: ``-1``. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: array: The computed NLL loss. """ def gaussian_nll_loss( inputs: mx.array, targets: mx.array, vars: mx.array, full: bool = ..., eps: float = ..., reduction: Reduction = ..., ) -> mx.array: r""" Computes the negative log likelihood loss for a Gaussian distribution. The loss is given by: .. math:: \frac{1}{2}\left(\log\left(\max\left(\text{vars}, \ \epsilon\right)\right) + \frac{\left(\text{inputs} - \text{targets} \right)^2} {\max\left(\text{vars}, \ \epsilon \right)}\right) + \text{const.} where ``inputs`` are the predicted means and ``vars`` are the the predicted variances. Args: inputs (array): The predicted expectation of the Gaussian distribution. targets (array): The target values (samples from the Gaussian distribution). vars (array): The predicted variance of the Gaussian distribution. full (bool, optional): Whether to include the constant term in the loss calculation. Default: ``False``. eps (float, optional): Small positive constant for numerical stability. Default: ``1e-6``. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: array: The Gaussian NLL loss. """ def kl_div_loss( inputs: mx.array, targets: mx.array, axis: int = ..., reduction: Reduction = ... ) -> mx.array: """ Computes the Kullback-Leibler divergence loss. Computes the following when ``reduction == 'none'``: .. code-block:: python mx.exp(targets) * (targets - inputs).sum(axis) Args: inputs (array): Log probabilities for the predicted distribution. targets (array): Log probabilities for the target distribution. axis (int, optional): The distribution axis. Default: ``-1``. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: array: The computed Kullback-Leibler divergence loss. """ def smooth_l1_loss( predictions: mx.array, targets: mx.array, beta: float = ..., reduction: Reduction = ..., ) -> mx.array: r""" Computes the smooth L1 loss. The smooth L1 loss is a variant of the L1 loss which replaces the absolute difference with a squared difference when the absolute difference is less than ``beta``. The formula for the smooth L1 Loss is: .. math:: l = \begin{cases} 0.5 (x - y)^2 / \beta, & \text{if } |x - y| < \beta \\ |x - y| - 0.5 \beta, & \text{otherwise} \end{cases} Args: predictions (array): Predicted values. targets (array): Ground truth values. beta (float, optional): The threshold after which the loss changes from the squared to the absolute difference. Default: ``1.0``. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``. Returns: array: The computed smooth L1 loss. """ def triplet_loss( anchors: mx.array, positives: mx.array, negatives: mx.array, axis: int = ..., p: int = ..., margin: float = ..., eps: float = ..., reduction: Reduction = ..., ) -> mx.array: r""" Computes the triplet loss for a set of anchor, positive, and negative samples. Margin is represented with alpha in the math section. .. math:: \max\left(\|A - P\|_p - \|A - N\|_p + \alpha, 0\right) Args: anchors (array): The anchor samples. positives (array): The positive samples. negatives (array): The negative samples. axis (int, optional): The distribution axis. Default: ``-1``. p (int, optional): The norm degree for pairwise distance. Default: ``2``. margin (float, optional): Margin for the triplet loss. Defaults to ``1.0``. eps (float, optional): Small positive constant to prevent numerical instability. Defaults to ``1e-6``. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: array: Computed triplet loss. If reduction is "none", returns a tensor of the same shape as input; if reduction is "mean" or "sum", returns a scalar tensor. """ def hinge_loss( inputs: mx.array, targets: mx.array, reduction: Reduction = ... ) -> mx.array: r""" Computes the hinge loss between inputs and targets. .. math:: \text{hinge}(y, y_{\text{pred}}) = \max(0, 1 - y \cdot y_{\text{pred}}) Args: inputs (array): The predicted values. targets (array): The target values. They should be -1 or 1. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: array: The computed hinge loss. """ def huber_loss( inputs: mx.array, targets: mx.array, delta: float = ..., reduction: Reduction = ... ) -> mx.array: r""" Computes the Huber loss between inputs and targets. .. math:: l_{\delta}(a) = \left\{ \begin{array}{ll} \frac{1}{2} a^2 & \text{for } |a| \leq \delta, \\ \delta \left( |a| - \frac{1}{2} \delta \right) & \text{otherwise.} \end{array} \right. Args: inputs (array): The predicted values. targets (array): The target values. delta (float, optional): The threshold at which to change between L1 and L2 loss. Default: ``1.0``. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: array: The computed Huber loss. """ def log_cosh_loss( inputs: mx.array, targets: mx.array, reduction: Reduction = ... ) -> mx.array: r""" Computes the log cosh loss between inputs and targets. Logcosh acts like L2 loss for small errors, ensuring stable gradients, and like the L1 loss for large errors, reducing sensitivity to outliers. This dual behavior offers a balanced, robust approach for regression tasks. .. math:: \text{logcosh}(y_{\text{true}}, y_{\text{pred}}) = \frac{1}{n} \sum_{i=1}^{n} \log(\cosh(y_{\text{pred}}^{(i)} - y_{\text{true}}^{(i)})) Args: inputs (array): The predicted values. targets (array): The target values. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: array: The computed log cosh loss. """ def cosine_similarity_loss( x1: mx.array, x2: mx.array, axis: int = ..., eps: float = ..., reduction: Reduction = ..., ) -> mx.array: r""" Computes the cosine similarity between the two inputs. The cosine similarity loss is given by .. math:: \frac{x_1 \cdot x_2}{\max(\|x_1\| \cdot \|x_2\|, \epsilon)} Args: x1 (mx.array): The first set of inputs. x2 (mx.array): The second set of inputs. axis (int, optional): The embedding axis. Default: ``1``. eps (float, optional): The minimum value of the denominator used for numerical stability. Default: ``1e-8``. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: mx.array: The computed cosine similarity loss. """ def margin_ranking_loss( inputs1: mx.array, inputs2: mx.array, targets: mx.array, margin: float = ..., reduction: Reduction = ..., ) -> mx.array: r""" Calculate the margin ranking loss that loss given inputs :math:`x_1`, :math:`x_2` and a label :math:`y` (containing 1 or -1). The loss is given by: .. math:: \text{loss} = \max (0, -y * (x_1 - x_2) + \text{margin}) Where :math:`y` represents ``targets``, :math:`x_1` represents ``inputs1`` and :math:`x_2` represents ``inputs2``. Args: inputs1 (array): Scores for the first input. inputs2 (array): Scores for the second input. targets (array): Labels indicating whether samples in ``inputs1`` should be ranked higher than samples in ``inputs2``. Values should be 1 or -1. margin (float, optional): The margin by which the scores should be separated. Default: ``0.0``. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: array: The computed margin ranking loss. Examples: >>> import mlx.core as mx >>> import mlx.nn as nn >>> targets = mx.array([1, 1, -1]) >>> inputs1 = mx.array([-0.573409, -0.765166, -0.0638]) >>> inputs2 = mx.array([0.75596, 0.225763, 0.256995]) >>> loss = nn.losses.margin_ranking_loss(inputs1, inputs2, targets) >>> loss array(0.773433, dtype=float32) """ ================================================ FILE: .mlx_typings/mlx/nn/utils.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Any, Callable, Optional import mlx.core as mx from .layers.base import Module def value_and_grad( model: Module, fn: Callable ): # -> _Wrapped[..., Any, ..., tuple[Any, Any]]: """Transform the passed function ``fn`` to a function that computes the gradients of ``fn`` wrt the model's trainable parameters and also its value. Args: model (Module): The model whose trainable parameters to compute gradients for fn (Callable): The scalar function to compute gradients for Returns: A callable that returns the value of ``fn`` and the gradients wrt the trainable parameters of ``model`` """ def checkpoint( module: Module, fn: Optional[Callable] = ... ): # -> _Wrapped[..., Any, ..., Any]: """Transform the passed callable to one that performs gradient checkpointing with respect to the trainable parameters of the module (and the callable's inputs). Args: module (Module): The module for whose parameters we will be performing gradient checkpointing. fn (Callable, optional): The function to checkpoint. If not provided it defaults to the provided module. Returns: A callable that saves the inputs and outputs during the forward pass and recomputes all intermediate states during the backward pass. """ def average_gradients( gradients: Any, group: Optional[mx.distributed.Group] = ..., all_reduce_size: int = ..., communication_type: Optional[mx.Dtype] = ..., communication_stream: Optional[mx.Stream] = ..., ): # -> Any: """Average the gradients across the distributed processes in the passed group. This helper enables concatenating several gradients of small arrays to one big all reduce call for better networking performance. Args: gradients (Any): The Python tree containing the gradients (it should have the same structure across processes) group (Optional[mlx.core.distributed.Group]): The group of processes to average the gradients. If set to ``None`` the global group is used. Default: ``None``. all_reduce_size (int): Group arrays until their size in bytes exceeds this number. Perform one communication step per group of arrays. If less or equal to 0 array grouping is disabled. Default: ``32MiB``. communication_type (Optional[mlx.core.Dtype]): If provided cast to this type before performing the communication. Typically cast to a smaller float to reduce the communication size. Default: ``None``. communication_stream (Optional[mlx.core.Stream]): The stream to usse for the communication. If unspecified the default communication stream is used which can vary by back-end. Default: ``None``. """ ================================================ FILE: .mlx_typings/mlx/utils.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Any, Callable, Dict, List, Optional, Tuple, Union from mlx.core import MX_ARRAY_TREE def tree_map( fn: Callable[..., Any], tree: Any, *rest: Any, is_leaf: Callable[..., bool] | None = ..., ) -> Any: """Applies ``fn`` to the leaves of the Python tree ``tree`` and returns a new collection with the results. If ``rest`` is provided, every item is assumed to be a superset of ``tree`` and the corresponding leaves are provided as extra positional arguments to ``fn``. In that respect, :meth:`tree_map` is closer to :func:`itertools.starmap` than to :func:`map`. The keyword argument ``is_leaf`` decides what constitutes a leaf from ``tree`` similar to :func:`tree_flatten`. .. code-block:: python import mlx.nn as nn from mlx.utils import tree_map model = nn.Linear(10, 10) print(model.parameters().keys()) # dict_keys(['weight', 'bias']) # square the parameters model.update(tree_map(lambda x: x*x, model.parameters())) Args: fn (callable): The function that processes the leaves of the tree. tree (Any): The main Python tree that will be iterated upon. rest (tuple[Any]): Extra trees to be iterated together with ``tree``. is_leaf (callable, optional): An optional callable that returns ``True`` if the passed object is considered a leaf or ``False`` otherwise. Returns: A Python tree with the new values returned by ``fn``. """ def tree_map_with_path( fn: Callable[..., Any], tree: Any, *rest: Any, is_leaf: Callable[..., bool] | None = ..., path: str | None = ..., ) -> Any: """Applies ``fn`` to the path and leaves of the Python tree ``tree`` and returns a new collection with the results. This function is the same :func:`tree_map` but the ``fn`` takes the path as the first argument followed by the remaining tree nodes. Args: fn (callable): The function that processes the leaves of the tree. tree (Any): The main Python tree that will be iterated upon. rest (tuple[Any]): Extra trees to be iterated together with ``tree``. is_leaf (Optional[Callable]): An optional callable that returns ``True`` if the passed object is considered a leaf or ``False`` otherwise. path (Optional[Any]): Prefix will be added to the result. Returns: A Python tree with the new values returned by ``fn``. Example: >>> from mlx.utils import tree_map_with_path >>> tree = {"model": [{"w": 0, "b": 1}, {"w": 0, "b": 1}]} >>> new_tree = tree_map_with_path(lambda path, _: print(path), tree) model.0.w model.0.b model.1.w model.1.b """ def tree_flatten( tree: Any, prefix: str = ..., is_leaf: Callable[..., bool] | None = ..., destination: list[tuple[str, Any]] | dict[str, Any] | None = ..., ) -> list[tuple[str, Any]] | dict[str, Any]: """Flattens a Python tree to a list of key, value tuples. The keys are using the dot notation to define trees of arbitrary depth and complexity. .. code-block:: python from mlx.utils import tree_flatten print(tree_flatten([[[0]]])) # [("0.0.0", 0)] print(tree_flatten([[[0]]], prefix=".hello")) # [("hello.0.0.0", 0)] tree_flatten({"a": {"b": 1}}, destination={}) {"a.b": 1} .. note:: Dictionaries should have keys that are valid Python identifiers. Args: tree (Any): The Python tree to be flattened. prefix (str): A prefix to use for the keys. The first character is always discarded. is_leaf (callable): An optional callable that returns True if the passed object is considered a leaf or False otherwise. destination (list or dict, optional): A list or dictionary to store the flattened tree. If None an empty list will be used. Default: ``None``. Returns: Union[List[Tuple[str, Any]], Dict[str, Any]]: The flat representation of the Python tree. """ def tree_unflatten(tree: list[tuple[str, Any]] | dict[str, Any]) -> Any: """Recreate a Python tree from its flat representation. .. code-block:: python from mlx.utils import tree_unflatten d = tree_unflatten([("hello.world", 42)]) print(d) # {"hello": {"world": 42}} d = tree_unflatten({"hello.world": 42}) print(d) # {"hello": {"world": 42}} Args: tree (list[tuple[str, Any]] or dict[str, Any]): The flat representation of a Python tree. For instance as returned by :meth:`tree_flatten`. Returns: A Python tree. """ def tree_reduce( fn: Callable[[Any, Any], Any], tree: list[MX_ARRAY_TREE] | tuple[MX_ARRAY_TREE, ...] | dict[str, MX_ARRAY_TREE], initializer=..., is_leaf=..., ) -> None: """Applies a reduction to the leaves of a Python tree. This function reduces Python trees into an accumulated result by applying the provided function ``fn`` to the leaves of the tree. Example: >>> from mlx.utils import tree_reduce >>> tree = {"a": [1, 2, 3], "b": [4, 5]} >>> tree_reduce(lambda acc, x: acc + x, tree, 0) 15 Args: fn (callable): The reducer function that takes two arguments (accumulator, current value) and returns the updated accumulator. tree (Any): The Python tree to reduce. It can be any nested combination of lists, tuples, or dictionaries. initializer (Any, optional): The initial value to start the reduction. If not provided, the first leaf value is used. is_leaf (callable, optional): A function to determine if an object is a leaf, returning ``True`` for leaf nodes and ``False`` otherwise. Returns: Any: The accumulated value. """ def tree_merge( tree_a, tree_b, merge_fn=... ): # -> dict[Any, Any] | list[Any] | tuple[Any, *tuple[Any, ...]] | tuple[Any, ...]: """Merge two Python trees in one containing the values of both. It can be thought of as a deep dict.update method. Args: tree_a (Any): The first Python tree. tree_b (Any): The second Python tree. merge_fn (callable, optional): A function to merge leaves. Returns: The Python tree containing the values of both ``tree_a`` and ``tree_b``. """ ================================================ FILE: .mlx_typings/mlx_lm/__init__.pyi ================================================ import models as models import tokenizer_utils as tokenizer_utils from generate import * ================================================ FILE: .mlx_typings/mlx_lm/convert.pyi ================================================ """ This type stub file was generated by pyright. """ import argparse from typing import Callable, Optional, Union import mlx.nn as nn def mixed_quant_predicate_builder( recipe: str, model: nn.Module, group_size: int = ... ) -> Callable[[str, nn.Module, dict], Union[bool, dict]]: ... QUANT_RECIPES = ... MODEL_CONVERSION_DTYPES = ... def convert( hf_path: str, mlx_path: str = ..., quantize: bool = ..., q_group_size: int = ..., q_bits: int = ..., q_mode: str = ..., dtype: Optional[str] = ..., upload_repo: str = ..., revision: Optional[str] = ..., dequantize: bool = ..., quant_predicate: Optional[ Union[Callable[[str, nn.Module, dict], Union[bool, dict]], str] ] = ..., trust_remote_code: bool = ..., ): # -> None: ... def configure_parser() -> argparse.ArgumentParser: """ Configures and returns the argument parser for the script. Returns: argparse.ArgumentParser: Configured argument parser. """ def main(): # -> None: ... if __name__ == "__main__": ... ================================================ FILE: .mlx_typings/mlx_lm/generate.pyi ================================================ """ This type stub file was generated by pyright. """ import contextlib from dataclasses import dataclass from typing import Any, Callable, Generator, List, Optional, Tuple, Union import mlx.core as mx import mlx.nn as nn from transformers import PreTrainedTokenizer from .tokenizer_utils import TokenizerWrapper DEFAULT_PROMPT = ... DEFAULT_MAX_TOKENS = ... DEFAULT_TEMP = ... DEFAULT_TOP_P = ... DEFAULT_MIN_P = ... DEFAULT_TOP_K = ... DEFAULT_XTC_PROBABILITY = ... DEFAULT_XTC_THRESHOLD = ... DEFAULT_MIN_TOKENS_TO_KEEP = ... DEFAULT_SEED = ... DEFAULT_MODEL = ... DEFAULT_QUANTIZED_KV_START = ... def str2bool(string): # -> bool: ... def setup_arg_parser(): # -> ArgumentParser: """Set up and return the argument parser.""" generation_stream = ... @contextlib.contextmanager def wired_limit( model: nn.Module, streams: Optional[List[mx.Stream]] = ... ): # -> Generator[None, Any, None]: """ A context manager to temporarily change the wired limit. Note, the wired limit should not be changed during an async eval. If an async eval could be running pass in the streams to synchronize with prior to exiting the context manager. """ @dataclass class GenerationResponse: """ The output of :func:`stream_generate`. Args: text (str): The next segment of decoded text. This can be an empty string. token (int): The next token. from_draft (bool): Whether the token was generated by the draft model. logprobs (mx.array): A vector of log probabilities. prompt_tokens (int): The number of tokens in the prompt. prompt_tps (float): The prompt processing tokens-per-second. generation_tokens (int): The number of generated tokens. generation_tps (float): The tokens-per-second for generation. peak_memory (float): The peak memory used so far in GB. finish_reason (str): The reason the response is being sent: "length", "stop" or `None` """ text: str token: int logprobs: mx.array from_draft: bool prompt_tokens: int prompt_tps: float generation_tokens: int generation_tps: float peak_memory: float finish_reason: Optional[str] = ... def maybe_quantize_kv_cache( prompt_cache: Any, quantized_kv_start: int | None, kv_group_size: int | None, kv_bits: int | None, ) -> None: ... def generate_step( prompt: mx.array, model: nn.Module, *, max_tokens: int = ..., sampler: Optional[Callable[[mx.array], mx.array]] = ..., logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = ..., max_kv_size: Optional[int] = ..., prompt_cache: Optional[Any] = ..., prefill_step_size: int = ..., kv_bits: Optional[int] = ..., kv_group_size: int = ..., quantized_kv_start: int = ..., prompt_progress_callback: Optional[Callable[[int], int]] = ..., input_embeddings: Optional[mx.array] = ..., ) -> Generator[Tuple[mx.array, mx.array], None, None]: """ A generator producing token ids based on the given prompt from the model. Args: prompt (mx.array): The input prompt. model (nn.Module): The model to use for generation. max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite generator. Default: ``256``. sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a token from a vector of log probabilities. Default: ``None``. logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional): A list of functions that take tokens and logits and return the processed logits. Default: ``None``. max_kv_size (int, optional): Maximum size of the key-value cache. Old entries (except the first 4 tokens) will be overwritten. prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if provided, the cache will be updated in place. prefill_step_size (int): Step size for processing the prompt. kv_bits (int, optional): Number of bits to use for KV cache quantization. None implies no cache quantization. Default: ``None``. kv_group_size (int): Group size for KV cache quantization. Default: ``64``. quantized_kv_start (int): Step to begin using a quantized KV cache. when ``kv_bits`` is non-None. Default: ``0``. prompt_progress_callback (Callable[[int], int]): A call-back which takes the prompt tokens processed so far and the total number of prompt tokens. input_embeddings (mx.array, optional): Input embeddings to use instead of or in conjunction with prompt tokens. Default: ``None``. Yields: Tuple[mx.array, mx.array]: One token and a vector of log probabilities. """ def speculative_generate_step( prompt: mx.array, model: nn.Module, draft_model: nn.Module, *, num_draft_tokens: int = ..., max_tokens: int = ..., sampler: Optional[Callable[[mx.array], mx.array]] = ..., logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = ..., prompt_cache: Optional[Any] = ..., prefill_step_size: int = ..., kv_bits: Optional[int] = ..., kv_group_size: int = ..., quantized_kv_start: int = ..., ) -> Generator[Tuple[mx.array, mx.array, bool], None, None]: """ A generator producing token ids based on the given prompt from the model. Args: prompt (mx.array): The input prompt. model (nn.Module): The model to use for generation. draft_model (nn.Module): The draft model for speculative decoding. num_draft_tokens (int, optional): The number of draft tokens for speculative decoding. Default: ``2``. max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite generator. Default: ``256``. sampler (Callable[[mx.array], mx.array], optional): A sampler for sampling a token from a vector of log probabilities. Default: ``None``. logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional): A list of functions that take tokens and logits and return the processed logits. Default: ``None``. prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if provided, the cache will be updated in place. The cache must be trimmable. prefill_step_size (int): Step size for processing the prompt. kv_bits (int, optional): Number of bits to use for KV cache quantization. None implies no cache quantization. Default: ``None``. kv_group_size (int): Group size for KV cache quantization. Default: ``64``. quantized_kv_start (int): Step to begin using a quantized KV cache. when ``kv_bits`` is non-None. Default: ``0``. Yields: Tuple[mx.array, mx.array, bool]: One token, a vector of log probabilities, and a bool indicating if the token was generated by the draft model """ def stream_generate( model: nn.Module, tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper], prompt: Union[str, mx.array, List[int]], max_tokens: int = ..., draft_model: Optional[nn.Module] = ..., **kwargs: object, ) -> Generator[GenerationResponse, None, None]: """ A generator producing text based on the given prompt from the model. Args: model (nn.Module): The model to use for generation. tokenizer (PreTrainedTokenizer): The tokenizer. prompt (Union[str, mx.array, List[int]]): The input prompt string or integer tokens. max_tokens (int): The maximum number of tokens to generate. Default: ``256``. draft_model (Optional[nn.Module]): An optional draft model. If provided then speculative decoding is used. The draft model must use the same tokenizer as the main model. Default: ``None``. kwargs: The remaining options get passed to :func:`generate_step`. See :func:`generate_step` for more details. Yields: GenerationResponse: An instance containing the generated text segment and associated metadata. See :class:`GenerationResponse` for details. """ def generate( model: nn.Module, tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper], prompt: Union[str, List[int]], verbose: bool = ..., **kwargs, ) -> str: """ Generate a complete response from the model. Args: model (nn.Module): The language model. tokenizer (PreTrainedTokenizer): The tokenizer. prompt (Union[str, List[int]]): The input prompt string or integer tokens. verbose (bool): If ``True``, print tokens and timing information. Default: ``False``. kwargs: The remaining options get passed to :func:`stream_generate`. See :func:`stream_generate` for more details. """ @dataclass class BatchStats: """ An data object to hold generation stats. Args: prompt_tokens (int): The number of prompt tokens processed. prompt_tps (float): The prompt processing tokens-per-second. prompt_time (float): The time in seconds spent in prompt processing. generation_tokens (int): The number of generated tokens. generation_tps (float): The tokens-per-second for generation. generation_time (float): The time in seconds spent in generation . peak_memory (float): The peak memory used so far in GB. """ prompt_tokens: int = ... prompt_tps: float = ... prompt_time: float = ... generation_tokens: int = ... generation_tps: float = ... generation_time: float = ... peak_memory: float = ... @dataclass class BatchResponse: """ An data object to hold a batch generation response. Args: texts: (List[str]): The generated text for each prompt. stats (BatchStats): Statistics about the generation. """ texts: List[str] stats: BatchStats caches: Optional[List[List[Any]]] def _left_pad_prompts(prompts: Any, max_length: Optional[int] = ...) -> mx.array: ... def _right_pad_prompts(prompts: Any, max_length: Optional[int] = ...) -> mx.array: ... def _make_cache( model: Any, left_padding: Any, max_kv_size: Optional[int] ) -> List[Any]: ... def _merge_caches(caches: Any) -> List[Any]: ... @dataclass class Batch: uids: List[int] y: mx.array logprobs: mx.array max_tokens: List[int] num_tokens: List[int] cache: List[Any] samplers: List[Any] logits_processors: List[Any] tokens: List[mx.array] def __len__(self) -> int: ... def filter(self, keep_idx: List[int]) -> None: ... def extend(self, other: "Batch") -> None: ... def extract_cache(self, idx: int) -> List[Any]: ... class BatchGenerator: model: Any max_kv_size: Optional[int] prefill_step_size: int unprocessed_prompts: List[Any] active_batch: Optional[Batch] prompt_progress_callback: Callable[[List[Tuple[int, int, int]]], None] _stats: BatchStats @dataclass class Response: uid: int token: int logprobs: mx.array finish_reason: Optional[str] prompt_cache: Any def __init__( self, model: Any, max_tokens: int = ..., stop_tokens: Optional[set[int]] = ..., sampler: Optional[Callable[[mx.array], mx.array]] = ..., logits_processors: Optional[ List[Callable[[mx.array, mx.array], mx.array]] ] = ..., completion_batch_size: int = ..., prefill_batch_size: int = ..., prefill_step_size: int = ..., prompt_progress_callback: Optional[ Callable[[List[Tuple[int, int, int]]], None] ] = ..., max_kv_size: Optional[int] = ..., ) -> None: ... def close(self) -> None: ... def insert( self, prompts: Any, max_tokens: Union[List[int], int, None] = ..., caches: Any = ..., samplers: Optional[List[Any]] = ..., logits_processors: Optional[List[Any]] = ..., ) -> List[int]: ... def remove( self, uids: List[int], return_prompt_caches: bool = ... ) -> Optional[dict[int, List[Any]]]: ... def stats(self) -> BatchStats: ... def next(self) -> List[Response]: ... def _process_prompts(self, prompts: List[Any]) -> Batch: ... def _step( self, input_tokens: mx.array, prompt_cache: List[Any], samplers: Optional[List[Any]], logits_processors: Optional[List[Any]], tokens: List[mx.array], ) -> Tuple[mx.array, List[mx.array]]: ... def batch_generate( model, tokenizer, prompts: List[int], max_tokens: Union[int, List[int]] = ..., verbose: bool = ..., **kwargs, ) -> BatchResponse: """ Generate responses for the given batch of prompts. Args: model (nn.Module): The language model. tokenizer (PreTrainedTokenizer): The tokenizer. prompt (List[List[int]]): The input prompts. verbose (bool): If ``True``, print tokens and timing information. Default: ``False``. max_tokens (Union[int, List[int]): Maximum number of output tokens. This can be per prompt if a list is provided. kwargs: The remaining options get passed to :obj:`BatchGenerator`. See :obj:`BatchGenerator` for more details. """ def main(): # -> None: ... if __name__ == "__main__": ... ================================================ FILE: .mlx_typings/mlx_lm/models/__init__.pyi ================================================ import cache as cache ================================================ FILE: .mlx_typings/mlx_lm/models/base.pyi ================================================ """ This type stub file was generated by pyright. """ from dataclasses import dataclass from typing import Optional import mlx.core as mx @dataclass class BaseModelArgs: @classmethod def from_dict(cls, params): # -> Self: ... def create_causal_mask( N: int, offset: int = ..., window_size: Optional[int] = ..., right_padding: Optional[mx.array] = ..., left_padding: Optional[mx.array] = ..., ): # -> array: ... def create_attention_mask( h, cache=..., window_size: Optional[int] = ..., return_array: bool = ... ): # -> array | Literal['causal'] | None: ... def create_ssm_mask(h, cache=...): # -> None: ... def quantized_scaled_dot_product_attention( queries: mx.array, q_keys: tuple[mx.array, mx.array, mx.array], q_values: tuple[mx.array, mx.array, mx.array], scale: float, mask: Optional[mx.array], group_size: int = ..., bits: int = ..., ) -> mx.array: ... def scaled_dot_product_attention( queries, keys, values, cache, scale: float, mask: Optional[mx.array], sinks: Optional[mx.array] = ..., ) -> mx.array: ... ================================================ FILE: .mlx_typings/mlx_lm/models/bitlinear_layers.pyi ================================================ """ This type stub file was generated by pyright. """ import mlx.nn as nn def bitnet_quantize(model, quantization_config: dict): ... def make_bitlinear_kernel(): """ Custom Metal kernel that performs matrix multiplication directly on packed weights and scales the output. This eliminates the need to store unpacked weights in memory. """ _bitlinear_kernel = ... class BitLinear(nn.Module): """ BitLinear module with memory-efficient weight handling. """ def __init__( self, in_features, out_features, bias=..., invert_weight_scales=... ) -> None: ... def execute_matmul_kernel(self, x, packed_weights): ... def __call__(self, x): # -> array: ... ================================================ FILE: .mlx_typings/mlx_lm/models/cache.pyi ================================================ """ This type stub file was generated by pyright. """ from typing import Any, Dict, List, Optional, Protocol, Literal, Self import mlx.nn as nn from mlx.core import array import mlx.core as mx class Cache(Protocol): keys: mx.array values: mx.array offset: int def update_and_fetch( self, keys: mx.array, values: mx.array ) -> tuple[mx.array, mx.array]: ... @property def state(self) -> tuple[mx.array | None, mx.array | None]: ... @state.setter def state(self, v) -> None: ... def make_prompt_cache( model: nn.Module, max_kv_size: Optional[int] = ... ) -> List[Cache | Any]: """ Construct the model's cache for use in generation. This function will defer the cache construction to the model if it has a ``make_cache`` method, otherwise it will make a default KV cache. Args: model (nn.Module): The language model. max_kv_size (Optional[int]): If provided and the model does not have a ``make_cache`` method, a ``RotatingKVCache`` is used with a maximum size of ``max_kv_size`` """ def save_prompt_cache( file_name: str, cache: List[Cache], metadata: Dict[str, str] = ... ) -> None: """ Save a pre-computed prompt cache to a file. Args: file_name (str): The ``.safetensors`` file name. cache (List[Any]): The model state. metadata (Dict[str, str]): Optional metadata to save along with model state. """ def load_prompt_cache(file_name: str, return_metadata=...) -> array: """ Load a prompt cache from a file. Args: file_name (str): The ``.safetensors`` file name. return_metadata (bool): Whether or not to return metadata. Default: ``False``. Returns: List[Any] or Tuple[List[Any], Dict[str, str]]: The prompt cache and the metadata if requested. """ def can_trim_prompt_cache(cache: List[Cache]) -> bool: """ Check if model's cache can be trimmed. """ def trim_prompt_cache(cache: List[Cache], num_tokens: int) -> List[Cache]: """ Trim the model's cache by the given number of tokens. This function will trim the cache if possible (in-place) and return the number of tokens that were trimmed. Args: cache (List[Any]): The model's cache. num_tokens (int): The number of tokens to trim. Returns: (int): The number of tokens that were trimmed. """ def create_attention_mask( N: int, offset: int, return_array: bool, window_size: Optional[int] ) -> array | Literal["causal"] | None: ... class _BaseCache(Cache): keys: mx.array values: mx.array offset: int @property def state(self) -> tuple[mx.array | None, mx.array | None]: ... @state.setter def state(self, v) -> None: ... @property def meta_state(self) -> Literal[""]: ... @meta_state.setter def meta_state(self, v) -> None: ... def trim(self, n: int) -> int: ... def is_trimmable(self) -> Literal[False]: ... @classmethod def from_state(cls, state, meta_state) -> Self: ... class ConcatenateKVCache(_BaseCache): """ConcatenateKVCache the simplest KV cache implementation. Can be used as a mock KV cache or when large blocks are being processed at a time in which case KVCache isn't necessarily faster. Consider using the KVCache with a larger step size before using this cache. """ def __init__(self) -> None: ... def update_and_fetch(self, keys, values): # -> tuple[Any | array, Any | array]: ... @property def state(self) -> tuple[mx.array | None, mx.array | None]: ... @state.setter def state(self, v): # -> None: ... def is_trimmable(self): # -> Literal[True]: ... def trim(self, n: int) -> int: ... def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None: ... class QuantizedKVCache(_BaseCache): step = ... def __init__(self, group_size: int = ..., bits: int = ...) -> None: ... def update_and_fetch(self, keys, values): # -> Any: ... @property def state(self) -> tuple[mx.array | None, mx.array | None]: ... @state.setter def state(self, v): # -> None: ... @property def meta_state(self): # -> tuple[str, ...]: ... @meta_state.setter def meta_state(self, v): # -> None: ... def is_trimmable(self): # -> Literal[True]: ... def trim(self, n: int) -> int: ... def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None: ... class KVCache(_BaseCache): step = ... def __init__(self) -> None: ... def update_and_fetch(self, keys, values): # -> tuple[array | Any, array | Any]: ... @property def state( self, ) -> tuple[mx.array | None, mx.array | None]: ... @state.setter def state(self, v) -> None: ... def is_trimmable(self): # -> Literal[True]: ... def trim(self, n: int) -> int: ... def to_quantized( self, group_size: int = ..., bits: int = ... ) -> QuantizedKVCache: ... def make_mask( self, *args: Any, **kwargs: Any ) -> mx.array | Literal["causal"] | None: ... class RotatingKVCache(_BaseCache): step = ... keys: mx.array | None values: mx.array | None keep: int max_size: int _idx: int def __init__(self, max_size, keep=...) -> None: ... def _trim( self, trim_size: int, v: mx.array, append: mx.array | None = ... ) -> mx.array: ... def update_and_fetch( self, keys, values ): # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]: ... @property def state( self, ) -> tuple[mx.array | None, mx.array | None]: ... @state.setter def state(self, v): # -> None: ... @property def meta_state(self): # -> tuple[str, ...]: ... @meta_state.setter def meta_state(self, v): # -> None: ... def is_trimmable(self): # -> bool: ... def trim(self, n: int) -> int: ... def to_quantized( self, group_size: int = ..., bits: int = ... ) -> QuantizedKVCache: ... def make_mask( self, N: int, window_size: Optional[int] = ..., return_array: bool = ... ): # -> array | Literal['causal'] | None: ... class ArraysCache(_BaseCache): def __init__(self, size, left_padding: Optional[List[int]] = ...) -> None: ... def __setitem__(self, idx, value): # -> None: ... def __getitem__(self, idx): ... @property def state(self) -> tuple[mx.array | None, mx.array | None]: ... @state.setter def state(self, v): # -> None: ... def filter(self, batch_indices): # -> None: """ In-place filter to keep just the given indices in the cache. """ def extend(self, other): # -> None: """ In-place extend this cache with the other cache. """ def make_mask(self, N: int) -> mx.array | None: ... class MambaCache(ArraysCache): def __init__(self, left_padding: Optional[List[int]] = ...) -> None: ... class ChunkedKVCache(KVCache): def __init__(self, chunk_size) -> None: ... def maybe_trim_front(self): # -> None: ... def update_and_fetch(self, keys, values): # -> tuple[array, array]: ... def trim(self, n: int) -> int: ... @property def meta_state(self): # -> tuple[str, ...]: ... @meta_state.setter def meta_state(self, v): # -> None: ... class CacheList(_BaseCache): def __init__(self, *caches) -> None: ... def __getitem__(self, idx): ... def is_trimmable(self): # -> bool: ... def trim(self, n: int) -> int: ... @property def state(self) -> list[tuple[mx.array | None, mx.array | None]]: ... @state.setter def state(self, v): # -> None: ... def filter(self, batch_indices): # -> None: """ In-place filter to keep just the given indices in the cache. """ def extend(self, other): # -> None: """ In-place extend this cache with the other cache. """ class BatchKVCache(_BaseCache): step = ... def __init__(self, left_padding: List[int]) -> None: """ The BatchKV cache expects inputs to be left-padded. E.g. the following prompts: [1, 3, 5] [7] [2, 6, 8, 9] Should be padded like so: [0, 1, 3, 5] [0, 0, 0, 7] [2, 6, 8, 9] And ``left_padding`` specifies the amount of padding for each. In this case, ``left_padding = [1, 3, 0]``. """ def update_and_fetch(self, keys, values): # -> tuple[array | Any, array | Any]: ... @property def state( self, ): # -> tuple[Any | array | None, Any | array | None, array | Any, array | Any]: ... @state.setter def state(self, v): # -> None: ... def is_trimmable(self): # -> Literal[True]: ... def trim(self, n): # -> int | float: ... def make_mask(self, N: int, return_array: bool = ..., **kwargs): # -> array: ... def filter(self, batch_indices): # -> None: """ In-place filter to keep just the given indices in the cache. """ def extend(self, other): # -> None: """ In-place extend this cache with the other cache. """ class BatchRotatingKVCache(_BaseCache): step = ... def __init__(self, max_size, left_padding: List[int]) -> None: ... def update_and_fetch( self, keys, values ): # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]: ... @property def state( self, ): # -> tuple[Any | array | None, Any | array | None, array | Any, array | Any]: ... @state.setter def state(self, v): # -> None: ... @property def meta_state(self): # -> tuple[str, ...]: ... @meta_state.setter def meta_state(self, v): # -> None: ... def is_trimmable(self): # -> bool: ... def trim(self, n): # -> int: ... def to_quantized( self, group_size: int = ..., bits: int = ... ) -> QuantizedKVCache: ... def make_mask( self, N: int, window_size: Optional[int] = ..., return_array: bool = ... ): # -> array: ... def filter(self, batch_indices): # -> None: """ In-place filter to keep just the given indices in the cache. """ def extend(self, other): # -> None: """ In-place extend this cache with the other cache. """ ================================================ FILE: .mlx_typings/mlx_lm/models/deepseek_v3.pyi ================================================ """Type stubs for mlx_lm.models.deepseek_v3""" from dataclasses import dataclass from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn from mlx_lm.models.mla import MultiLinear from .base import BaseModelArgs from .switch_layers import SwitchGLU @dataclass class ModelArgs(BaseModelArgs): model_type: str vocab_size: int hidden_size: int intermediate_size: int moe_intermediate_size: int num_hidden_layers: int num_attention_heads: int num_key_value_heads: int n_shared_experts: Optional[int] n_routed_experts: Optional[int] routed_scaling_factor: float kv_lora_rank: int q_lora_rank: Optional[int] qk_rope_head_dim: int v_head_dim: int qk_nope_head_dim: int topk_method: str scoring_func: str norm_topk_prob: bool n_group: int topk_group: int num_experts_per_tok: int moe_layer_freq: int first_k_dense_replace: int max_position_embeddings: int rms_norm_eps: float rope_theta: float rope_scaling: Optional[Dict[str, Any]] attention_bias: bool class DeepseekV3Attention(nn.Module): config: ModelArgs hidden_size: int num_heads: int max_position_embeddings: int rope_theta: float q_lora_rank: Optional[int] qk_rope_head_dim: int kv_lora_rank: int v_head_dim: int qk_nope_head_dim: int q_head_dim: int scale: float q_proj: nn.Linear q_a_proj: nn.Linear q_a_layernorm: nn.RMSNorm q_b_proj: nn.Linear kv_a_proj_with_mqa: nn.Linear kv_a_layernorm: nn.RMSNorm # kv_b_proj: nn.Linear embed_q: MultiLinear unembed_out: MultiLinear o_proj: nn.Linear rope: Any def __init__(self, config: ModelArgs) -> None: ... def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: ... class DeepseekV3MLP(nn.Module): config: ModelArgs hidden_size: int intermediate_size: int gate_proj: nn.Linear up_proj: nn.Linear down_proj: nn.Linear def __init__( self, config: ModelArgs, hidden_size: Optional[int] = None, intermediate_size: Optional[int] = None, ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class MoEGate(nn.Module): config: ModelArgs top_k: int norm_topk_prob: bool n_routed_experts: Optional[int] routed_scaling_factor: float n_group: int topk_group: int weight: mx.array e_score_correction_bias: mx.array def __init__(self, config: ModelArgs) -> None: ... def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ... class DeepseekV3MoE(nn.Module): config: ModelArgs num_experts_per_tok: int switch_mlp: SwitchGLU gate: MoEGate shared_experts: DeepseekV3MLP sharding_group: Optional[mx.distributed.Group] def __init__(self, config: ModelArgs) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class DeepseekV3DecoderLayer(nn.Module): self_attn: DeepseekV3Attention mlp: DeepseekV3MLP | DeepseekV3MoE input_layernorm: nn.RMSNorm post_attention_layernorm: nn.RMSNorm def __init__(self, config: ModelArgs, layer_idx: int) -> None: ... def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: ... class DeepseekV3Model(nn.Module): vocab_size: int embed_tokens: nn.Embedding layers: list[DeepseekV3DecoderLayer] norm: nn.RMSNorm def __init__(self, config: ModelArgs) -> None: ... def __call__( self, x: mx.array, cache: Optional[Any] = None, ) -> mx.array: ... class Model(nn.Module): model_type: str model: DeepseekV3Model lm_head: nn.Linear def __init__(self, config: ModelArgs) -> None: ... def __call__( self, inputs: mx.array, cache: Optional[Any] = None, ) -> mx.array: ... def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ... @property def layers(self) -> list[DeepseekV3DecoderLayer]: ... ================================================ FILE: .mlx_typings/mlx_lm/models/glm4_moe.pyi ================================================ from dataclasses import dataclass from typing import Any, Dict, Optional import mlx.core as mx import mlx.nn as nn from .base import BaseModelArgs from .switch_layers import SwitchGLU @dataclass class ModelArgs(BaseModelArgs): model_type: str vocab_size: int hidden_size: int intermediate_size: int max_position_embeddings: int moe_intermediate_size: int norm_topk_prob: bool num_attention_heads: int n_group: int head_dim: int topk_group: int n_shared_experts: int n_routed_experts: int routed_scaling_factor: float num_experts_per_tok: int first_k_dense_replace: int num_hidden_layers: int num_key_value_heads: int rms_norm_eps: float rope_theta: float rope_scaling: Optional[Dict[str, Any]] use_qk_norm: bool tie_word_embeddings: bool attention_bias: bool partial_rotary_factor: float scoring_func: str topk_method: str class Attention(nn.Module): n_heads: int n_kv_heads: int scale: float q_proj: nn.Linear k_proj: nn.Linear v_proj: nn.Linear o_proj: nn.Linear use_qk_norm: bool q_norm: nn.RMSNorm k_norm: nn.RMSNorm rope: nn.RoPE def __init__(self, args: ModelArgs) -> None: ... def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: ... class MLP(nn.Module): config: ModelArgs hidden_size: int intermediate_size: int gate_proj: nn.Linear up_proj: nn.Linear down_proj: nn.Linear def __init__( self, config: ModelArgs, hidden_size: Optional[int] = None, intermediate_size: Optional[int] = None, ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class MoEGate(nn.Module): config: ModelArgs top_k: int norm_topk_prob: bool n_routed_experts: int routed_scaling_factor: float n_group: int topk_group: int weight: mx.array e_score_correction_bias: mx.array def __init__(self, config: ModelArgs) -> None: ... def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ... class MoE(nn.Module): config: ModelArgs num_experts_per_tok: int switch_mlp: SwitchGLU gate: MoEGate shared_experts: MLP sharding_group: Optional[mx.distributed.Group] def __init__(self, config: ModelArgs) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class DecoderLayer(nn.Module): self_attn: Attention mlp: MLP | MoE input_layernorm: nn.RMSNorm post_attention_layernorm: nn.RMSNorm def __init__(self, config: ModelArgs, layer_idx: int) -> None: ... def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: ... class LanguageModel(nn.Module): vocab_size: int embed_tokens: nn.Embedding layers: list[DecoderLayer] norm: nn.RMSNorm pipeline_rank: int pipeline_size: int start_idx: int end_idx: Optional[int] num_layers: int def __init__(self, config: ModelArgs) -> None: ... def __call__( self, x: mx.array, cache: Optional[Any] = None, ) -> mx.array: ... @property def pipeline_layers(self) -> list[DecoderLayer]: ... class Model(nn.Module): args: ModelArgs model_type: str model: LanguageModel lm_head: nn.Linear def __init__(self, config: ModelArgs) -> None: ... def __call__( self, inputs: mx.array, cache: Optional[Any] = None, ) -> mx.array: ... def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ... def shard(self, group: Optional[mx.distributed.Group] = None) -> None: ... @property def layers(self) -> list[DecoderLayer]: ... @property def cast_predicate(self) -> Any: ... ================================================ FILE: .mlx_typings/mlx_lm/models/glm_moe_dsa.pyi ================================================ """Type stubs for mlx_lm.models.glm_moe_dsa""" from dataclasses import dataclass from typing import Any, Dict, Optional from .base import BaseModelArgs from .deepseek_v32 import Model as DSV32Model @dataclass class ModelArgs(BaseModelArgs): model_type: str vocab_size: int hidden_size: int index_head_dim: int index_n_heads: int index_topk: int intermediate_size: int moe_intermediate_size: int num_hidden_layers: int num_attention_heads: int num_key_value_heads: int n_shared_experts: Optional[int] n_routed_experts: Optional[int] routed_scaling_factor: float kv_lora_rank: int q_lora_rank: int qk_rope_head_dim: int v_head_dim: int qk_nope_head_dim: int topk_method: str scoring_func: str norm_topk_prob: bool n_group: int topk_group: int num_experts_per_tok: int moe_layer_freq: int first_k_dense_replace: int max_position_embeddings: int rms_norm_eps: float rope_parameters: Dict[str, Any] attention_bias: bool rope_scaling: Dict[str, Any] | None rope_theta: float | None class Model(DSV32Model): def __init__(self, config: ModelArgs) -> None: ... ================================================ FILE: .mlx_typings/mlx_lm/models/nemotron_h.pyi ================================================ from dataclasses import dataclass from typing import Any, List, Optional, Tuple import mlx.core as mx import mlx.nn as nn from .cache import ArraysCache, KVCache from .switch_layers import SwitchMLP @dataclass class ModelArgs: model_type: str vocab_size: int hidden_size: int intermediate_size: int num_hidden_layers: int max_position_embeddings: int num_attention_heads: int num_key_value_heads: int attention_bias: bool mamba_num_heads: int mamba_head_dim: int mamba_proj_bias: bool ssm_state_size: int conv_kernel: int n_groups: int mlp_bias: bool layer_norm_epsilon: float use_bias: bool use_conv_bias: bool hybrid_override_pattern: List[str] head_dim: Optional[int] moe_intermediate_size: Optional[int] moe_shared_expert_intermediate_size: Optional[int] n_group: Optional[int] n_routed_experts: Optional[int] n_shared_experts: Optional[int] topk_group: Optional[int] num_experts_per_tok: Optional[int] norm_topk_prob: Optional[bool] routed_scaling_factor: Optional[float] time_step_limit: Optional[Tuple[float, float]] time_step_min: Optional[float] time_step_max: Optional[float] @classmethod def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ... def __post_init__(self) -> None: ... class NemotronHMamba2Mixer(nn.Module): num_heads: int hidden_size: int ssm_state_size: int conv_kernel_size: int intermediate_size: int n_groups: int head_dim: int conv_dim: int conv1d: nn.Conv1d in_proj: nn.Linear dt_bias: mx.array A_log: mx.array D: mx.array norm: nn.RMSNorm heads_per_group: int out_proj: nn.Linear def __init__(self, args: ModelArgs) -> None: ... def __call__( self, hidden_states: mx.array, mask: Optional[mx.array], cache: Optional[ArraysCache] = None, ) -> mx.array: ... class NemotronHAttention(nn.Module): hidden_size: int num_heads: int head_dim: int num_key_value_heads: int scale: float q_proj: nn.Linear k_proj: nn.Linear v_proj: nn.Linear o_proj: nn.Linear def __init__(self, args: ModelArgs) -> None: ... def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[KVCache] = None, ) -> mx.array: ... class NemotronHMLP(nn.Module): up_proj: nn.Linear down_proj: nn.Linear def __init__( self, args: ModelArgs, intermediate_size: Optional[int] = None ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class NemotronHMoE(nn.Module): num_experts_per_tok: int switch_mlp: SwitchMLP shared_experts: NemotronHMLP def __init__(self, config: ModelArgs) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class NemotronHBlock(nn.Module): block_type: str norm: nn.RMSNorm mixer: NemotronHMamba2Mixer | NemotronHAttention | NemotronHMLP | NemotronHMoE def __init__(self, args: ModelArgs, block_type: str) -> None: ... def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: ... class NemotronHModel(nn.Module): embeddings: nn.Embedding layers: list[NemotronHBlock] norm_f: nn.RMSNorm fa_idx: int ssm_idx: int def __init__(self, args: ModelArgs) -> None: ... def __call__( self, inputs: mx.array, cache: Optional[Any] = None, ) -> mx.array: ... class Model(nn.Module): args: ModelArgs backbone: NemotronHModel lm_head: nn.Linear model_type: str def __init__(self, args: ModelArgs) -> None: ... def __call__( self, inputs: mx.array, cache: Optional[Any] = None, ) -> mx.array: ... @property def layers(self) -> list[NemotronHBlock]: ... def make_cache(self) -> list[ArraysCache | KVCache]: ... def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ... ================================================ FILE: .mlx_typings/mlx_lm/models/qwen3_5.pyi ================================================ from dataclasses import dataclass from typing import Any, Optional import mlx.core as mx import mlx.nn as nn from .cache import ArraysCache, KVCache from .qwen3_next import ( Qwen3NextAttention as Attention, Qwen3NextMLP as MLP, Qwen3NextRMSNormGated as RMSNormGated, Qwen3NextSparseMoeBlock, ) SparseMoeBlock = Qwen3NextSparseMoeBlock from .switch_layers import SwitchGLU @dataclass class TextModelArgs: model_type: str hidden_size: int intermediate_size: int num_hidden_layers: int num_attention_heads: int rms_norm_eps: float vocab_size: int num_key_value_heads: int max_position_embeddings: int linear_num_value_heads: int linear_num_key_heads: int linear_key_head_dim: int linear_value_head_dim: int linear_conv_kernel_dim: int tie_word_embeddings: bool attention_bias: bool head_dim: Optional[int] full_attention_interval: int num_experts: int num_experts_per_tok: int decoder_sparse_step: int shared_expert_intermediate_size: int moe_intermediate_size: int norm_topk_prob: bool rope_parameters: Optional[dict[str, Any]] partial_rotary_factor: float rope_theta: float rope_scaling: Optional[dict[str, Any]] @classmethod def from_dict(cls, params: dict[str, Any]) -> TextModelArgs: ... def __post_init__(self) -> None: ... class GatedDeltaNet(nn.Module): hidden_size: int num_v_heads: int num_k_heads: int head_k_dim: int head_v_dim: int key_dim: int value_dim: int conv_kernel_size: int conv_dim: int conv1d: nn.Conv1d in_proj_qkv: nn.Linear in_proj_z: nn.Linear in_proj_b: nn.Linear in_proj_a: nn.Linear dt_bias: mx.array A_log: mx.array norm: RMSNormGated out_proj: nn.Linear def __init__(self, config: TextModelArgs) -> None: ... def __call__( self, inputs: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: ... class DecoderLayer(nn.Module): is_linear: bool linear_attn: GatedDeltaNet self_attn: Attention input_layernorm: nn.RMSNorm post_attention_layernorm: nn.RMSNorm mlp: MLP | SparseMoeBlock def __init__(self, args: TextModelArgs, layer_idx: int) -> None: ... def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: ... class Qwen3_5TextModel(nn.Module): embed_tokens: nn.Embedding layers: list[DecoderLayer] norm: nn.RMSNorm ssm_idx: int fa_idx: int def __init__(self, args: TextModelArgs) -> None: ... def __call__( self, inputs: mx.array, cache: Optional[Any] = None, input_embeddings: Optional[mx.array] = None, ) -> mx.array: ... class TextModel(nn.Module): args: TextModelArgs model_type: str model: Qwen3_5TextModel lm_head: nn.Linear def __init__(self, args: TextModelArgs) -> None: ... def __call__( self, inputs: mx.array, cache: Optional[Any] = None, input_embeddings: Optional[mx.array] = None, ) -> mx.array: ... @property def layers(self) -> list[DecoderLayer]: ... def make_cache(self) -> list[ArraysCache | KVCache]: ... def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ... @dataclass class ModelArgs: model_type: str text_config: dict[str, Any] @classmethod def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ... class Model(nn.Module): args: ModelArgs model_type: str language_model: TextModel def __init__(self, args: ModelArgs) -> None: ... def __call__( self, inputs: mx.array, cache: Optional[Any] = None, input_embeddings: Optional[mx.array] = None, ) -> mx.array: ... def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ... @property def layers(self) -> list[DecoderLayer]: ... def make_cache(self) -> list[ArraysCache | KVCache]: ... ================================================ FILE: .mlx_typings/mlx_lm/models/qwen3_5_moe.pyi ================================================ from dataclasses import dataclass from typing import Any, Optional import mlx.core as mx import mlx.nn as nn from .cache import ArraysCache, KVCache from .qwen3_5 import DecoderLayer, Model as Qwen3_5Model, TextModel @dataclass class ModelArgs: model_type: str text_config: dict[str, Any] @classmethod def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ... class Model(Qwen3_5Model): def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ... ================================================ FILE: .mlx_typings/mlx_lm/models/qwen3_next.pyi ================================================ """Type stubs for mlx_lm.models.qwen3_next""" from typing import Any, Optional import mlx.core as mx import mlx.nn as nn from .cache import ArraysCache, KVCache from .switch_layers import SwitchGLU class Qwen3NextRMSNormGated(nn.Module): eps: float weight: mx.array def __init__(self, hidden_size: int, eps: float = ...) -> None: ... def __call__( self, hidden_states: mx.array, gate: mx.array | None = None ) -> mx.array: ... class Qwen3NextMLP(nn.Module): gate_proj: nn.Linear down_proj: nn.Linear up_proj: nn.Linear def __init__(self, dim: int, hidden_dim: int) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class Qwen3NextGatedDeltaNet(nn.Module): hidden_size: int num_v_heads: int num_k_heads: int head_k_dim: int head_v_dim: int key_dim: int value_dim: int conv_kernel_size: int conv_dim: int conv1d: nn.Conv1d in_proj_qkvz: nn.Linear in_proj_ba: nn.Linear dt_bias: mx.array A_log: mx.array out_proj: nn.Linear def __init__(self, config: Any) -> None: ... def __call__( self, inputs: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: ... class Qwen3NextAttention(nn.Module): num_attention_heads: int num_key_value_heads: int head_dim: int scale: float q_proj: nn.Linear k_proj: nn.Linear v_proj: nn.Linear o_proj: nn.Linear def __init__(self, args: Any) -> None: ... def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: ... class Qwen3NextSparseMoeBlock(nn.Module): norm_topk_prob: bool num_experts: int top_k: int gate: nn.Linear switch_mlp: SwitchGLU shared_expert: Qwen3NextMLP shared_expert_gate: nn.Linear def __init__(self, args: Any) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class Qwen3NextDecoderLayer(nn.Module): is_linear: bool linear_attn: Qwen3NextGatedDeltaNet self_attn: Qwen3NextAttention input_layernorm: nn.RMSNorm post_attention_layernorm: nn.RMSNorm mlp: Qwen3NextMLP | Qwen3NextSparseMoeBlock def __init__(self, args: Any, layer_idx: int) -> None: ... def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: ... class Qwen3NextModel(nn.Module): embed_tokens: nn.Embedding layers: list[Qwen3NextDecoderLayer] norm: nn.RMSNorm ssm_idx: int fa_idx: int def __init__(self, args: Any) -> None: ... def __call__( self, inputs: mx.array, cache: Optional[Any] = None, ) -> mx.array: ... class Model(nn.Module): model_type: str model: Qwen3NextModel lm_head: nn.Linear def __init__(self, args: Any) -> None: ... def __call__( self, inputs: mx.array, cache: Optional[Any] = None, ) -> mx.array: ... def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ... @property def layers(self) -> list[Qwen3NextDecoderLayer]: ... def make_cache(self) -> list[ArraysCache | KVCache]: ... ================================================ FILE: .mlx_typings/mlx_lm/models/step3p5.pyi ================================================ from dataclasses import dataclass from typing import Any, Dict, List, Optional import mlx.core as mx import mlx.nn as nn from .base import BaseModelArgs from .switch_layers import SwitchGLU @dataclass class ModelArgs(BaseModelArgs): model_type: str hidden_size: int num_hidden_layers: int vocab_size: int num_attention_heads: int num_attention_groups: int head_dim: int intermediate_size: int rms_norm_eps: float rope_theta: float rope_scaling: Optional[Dict[str, Any]] max_position_embeddings: int sliding_window: int layer_types: Optional[List[str]] yarn_only_types: Optional[List[str]] partial_rotary_factors: Optional[List[float]] attention_other_setting: Optional[Dict[str, Any]] use_head_wise_attn_gate: bool moe_num_experts: int moe_top_k: int moe_intermediate_size: int share_expert_dim: int moe_layers_enum: Optional[str] moe_router_scaling_factor: float norm_expert_weight: bool swiglu_limits: Optional[List[float]] swiglu_limits_shared: Optional[List[float]] tie_word_embeddings: bool class Step3p5MLP(nn.Module): hidden_size: int intermediate_size: int gate_proj: nn.Linear up_proj: nn.Linear down_proj: nn.Linear limit: Optional[float] def __init__( self, args: ModelArgs, intermediate_size: int, swiglu_limit: float = 0 ) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class Step3p5MoEGate(nn.Module): top_k: int n_routed_experts: int routed_scaling_factor: float norm_topk_prob: bool gate: nn.Linear router_bias: mx.array def __init__(self, args: ModelArgs) -> None: ... def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ... class Step3p5MoE(nn.Module): gate: Step3p5MoEGate switch_mlp: SwitchGLU share_expert: Step3p5MLP sharding_group: Optional[mx.distributed.Group] def __init__(self, args: ModelArgs, layer_idx: int) -> None: ... def __call__(self, x: mx.array) -> mx.array: ... class Step3p5Attention(nn.Module): is_sliding: bool num_heads: int num_kv_heads: int head_dim: int scale: float q_proj: nn.Linear k_proj: nn.Linear v_proj: nn.Linear o_proj: nn.Linear q_norm: nn.Module k_norm: nn.Module use_head_wise_attn_gate: bool g_proj: nn.Linear rope: nn.Module def __init__(self, args: ModelArgs, layer_idx: int) -> None: ... def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: ... class Step3p5DecoderLayer(nn.Module): self_attn: Step3p5Attention is_sliding: bool is_moe_layer: bool mlp: Step3p5MLP | Step3p5MoE input_layernorm: nn.Module post_attention_layernorm: nn.Module def __init__(self, args: ModelArgs, layer_idx: int) -> None: ... def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: ... class Step3p5Model(nn.Module): args: ModelArgs vocab_size: int num_layers: int embed_tokens: nn.Embedding layers: list[Step3p5DecoderLayer] norm: nn.Module _swa_idx: Optional[int] _full_idx: Optional[int] def __init__(self, args: ModelArgs) -> None: ... def __call__( self, x: mx.array, cache: Optional[List[Any]] = None, ) -> mx.array: ... class Model(nn.Module): args: ModelArgs model_type: str model: Step3p5Model lm_head: nn.Linear def __init__(self, args: ModelArgs) -> None: ... def __call__( self, inputs: mx.array, cache: Optional[List[Any]] = None, ) -> mx.array: ... def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ... def shard(self, group: Optional[mx.distributed.Group] = None) -> None: ... @property def layers(self) -> list[Step3p5DecoderLayer]: ... def make_cache(self) -> list[Any]: ... @property def cast_predicate(self) -> Any: ... @property def quant_predicate(self) -> Any: ... ================================================ FILE: .mlx_typings/mlx_lm/models/switch_layers.pyi ================================================ """ This type stub file was generated by pyright. """ from functools import partial import mlx.core as mx import mlx.nn as nn class QuantizedSwitchLinear(nn.Module): def __init__( self, input_dims: int, output_dims: int, num_experts: int, bias: bool = ..., group_size: int = ..., bits: int = ..., mode: str = ..., ) -> None: ... @property def input_dims(self): # -> int: ... @property def output_dims(self): # -> int: ... @property def num_experts(self): # -> int: ... def __call__(self, x, indices, sorted_indices=...): # -> array: ... class SwitchLinear(nn.Module): def __init__( self, input_dims: int, output_dims: int, num_experts: int, bias: bool = ... ) -> None: ... @property def input_dims(self): # -> int: ... @property def output_dims(self): # -> int: ... @property def num_experts(self): # -> int: ... def __call__(self, x, indices, sorted_indices=...): ... def to_quantized( self, group_size: int = ..., bits: int = ..., mode: str = ... ): # -> QuantizedSwitchLinear: ... @partial(mx.compile, shapeless=True) def swiglu(x, gate): ... class SwiGLU(nn.Module): def __init__(self) -> None: ... def __call__(self, x, gate): ... class SwitchGLU(nn.Module): gate_proj: SwitchLinear up_proj: SwitchLinear down_proj: SwitchLinear activation: SwiGLU def __init__( self, input_dims: int, hidden_dims: int, num_experts: int, activation=..., bias: bool = ..., ) -> None: ... def __call__(self, x, indices) -> mx.array: ... class SwitchMLP(nn.Module): fc1: SwitchLinear fc2: SwitchLinear def __init__( self, input_dims: int, hidden_dims: int, num_experts: int, activation=..., bias: bool = ..., ) -> None: ... def __call__(self, x, indices) -> mx.array: ... ================================================ FILE: .mlx_typings/mlx_lm/sample_utils.pyi ================================================ """ This type stub file was generated by pyright. """ from functools import partial from typing import Callable, Dict, List, Optional import mlx.core as mx def make_sampler( temp: float = ..., top_p: float = ..., min_p: float = ..., min_tokens_to_keep: int = ..., top_k: int = ..., xtc_probability: float = ..., xtc_threshold: float = ..., xtc_special_tokens: List[int] = ..., ) -> Callable[[mx.array], mx.array]: """ Make a sampler function for use with ``generate_step``. Args: temp (float): The temperature for sampling, if 0 the argmax is used. Default: ``0``. top_p (float, optional): Nulceus sampling, higher means model considers more less likely words. min_p (float, optional): The minimum value (scaled by the top token's probability) that a token probability must have to be considered. min_tokens_to_keep (int, optional): Minimum number of tokens that cannot be filtered by min_p sampling. top_k (int, optional): The top k tokens ranked by probability to constrain the sampling to. xtc_probability (float, optional): The probability of applying XTC sampling. xtc_threshold (float, optional): The threshold the probs need to reach for being sampled. xtc_special_tokens (list(int), optional): List of special tokens IDs to be excluded from XTC sampling. Returns: Callable[mx.array, mx.array]: A sampler which takes log-probabilities and returns tokens. """ def make_logits_processors( logit_bias: Optional[Dict[int, float]] = ..., repetition_penalty: Optional[float] = ..., repetition_context_size: Optional[int] = ..., ) -> list[Callable[[mx.array, mx.array], mx.array]]: """ Make logits processors for use with ``generate_step``. Args: repetition_penalty (float, optional): The penalty factor for repeating tokens. repetition_context_size (int, optional): The number of tokens to consider for repetition penalty. Default: ``20``. logit_bias (dictionary, optional): Additive logit bias. Returns: List[Callable[[mx.array, mx.array], mx.array]]: A list of logits processors. Each processor in the list is a callable which takes an array of tokens and an array of logits and returns the updated logits. """ @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state) def apply_top_k(logprobs: mx.array, top_k: int) -> mx.array: """ Sample from only the top K tokens ranked by probability. Args: logprobs: A vector of log probabilities. top_k (int): Top k tokens to sample from. """ @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state) def apply_min_p( logprobs: mx.array, min_p: float, min_tokens_to_keep: int = ... ) -> mx.array: """ Apply min-p sampling to the logprobs. Min-p keeps all tokens that are above a minimum probability, scaled by the probability of the most likely token. As a result, the filter is more aggressive given a very high-probability token. Args: logprobs: A vector of log probabilities. min_p (float): Minimum token probability. Typical values are in the 0.01-0.2 range, comparably selective as setting `top_p` in the 0.99-0.8 range. min_tokens_to_keep (int, optional): Minimum number of tokens that cannot be filtered. Default: ``1``. """ @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state) def apply_top_p(logprobs: mx.array, top_p: float) -> mx.array: """ Apply top-p (nucleus) sampling to logits. Args: logprobs: A vector of log probabilities. top_p: The cumulative probability threshold for top-p filtering. Returns: token selected based on the top-p criterion. """ @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state) def apply_xtc( logits: mx.array, xtc_probability: float, xtc_threshold: float, xtc_special_tokens: List[int], ) -> mx.array: """ Apply XTC sampling to the logits. Args: logits: The logits from the model's output. xtc_probability (float): Probability of XTC sampling to happen for each token xtc_threshold (float): The threshold the probs need to reach for being sampled. special_tokens_ids (list(int)): List of special tokens IDs to be excluded from XTC sampling. """ @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state) def categorical_sampling(logits, temp): # -> array: ... def make_repetition_penalty( penalty: float, context_size: int = ... ): # -> Callable[..., Any]: """ Make repetition penalty processor. Paper: https://arxiv.org/abs/1909.05858 Args: penalty (float): The repetition penalty factor to be applied. context_size (int): The number of previous tokens to use. Default: ``20``. Returns: Callable[[mx.array, List[int]], mx.array]: The repetition penalty processor. """ ================================================ FILE: .mlx_typings/mlx_lm/tokenizer_utils.pyi ================================================ """ This type stub file was generated by pyright. """ from functools import partial from pathlib import Path from typing import Any from transformers import PreTrainedTokenizerFast class StreamingDetokenizer: """The streaming detokenizer interface so that we can detokenize one token at a time. Example usage is as follows: detokenizer = ... # Reset the tokenizer state detokenizer.reset() for token in generate(...): detokenizer.add_token(token.item()) # Contains the whole text so far. Some tokens may not be included # since it contains whole words usually. detokenizer.text # Contains the printable segment (usually a word) since the last # time it was accessed detokenizer.last_segment # Contains all the tokens added so far detokenizer.tokens # Make sure that we detokenize any remaining tokens detokenizer.finalize() # Now detokenizer.text should match tokenizer.decode(detokenizer.tokens) """ __slots__ = ... def reset(self) -> None: ... def add_token(self, token: int) -> None: ... def finalize(self) -> None: ... @property def last_segment(self) -> str: """Return the last segment of readable text since last time this property was accessed.""" class NaiveStreamingDetokenizer(StreamingDetokenizer): """NaiveStreamingDetokenizer relies on the underlying tokenizer implementation and should work with every tokenizer. Its complexity is O(T^2) where T is the longest line since it will repeatedly detokenize the same tokens until a new line is generated. """ def __init__(self, tokenizer) -> None: ... def reset(self): # -> None: ... def add_token(self, token): # -> None: ... def finalize(self): # -> None: ... @property def text(self): # -> str: ... class SPMStreamingDetokenizer(StreamingDetokenizer): """A streaming detokenizer for SPM models. It adds tokens to the text if the next token starts with the special SPM underscore which results in linear complexity. """ def __init__(self, tokenizer, trim_space=...) -> None: ... def reset(self): # -> None: ... def add_token(self, token): # -> None: ... def finalize(self): # -> None: ... class BPEStreamingDetokenizer(StreamingDetokenizer): """A streaming detokenizer for OpenAI style BPE models. It adds tokens to the text if the next token starts with a space similar to the SPM detokenizer. """ _byte_decoder = ... _space_matches = ... def __init__(self, tokenizer) -> None: ... def reset(self): # -> None: ... def add_token(self, token): # -> None: ... def finalize(self): # -> None: ... @classmethod def make_byte_decoder(cls): # -> None: """See https://github.com/openai/gpt-2/blob/master/src/encoder.py for the rationale.""" class TokenizerWrapper: """A wrapper that combines an HF tokenizer and a detokenizer. Accessing any attribute other than the ``detokenizer`` is forwarded to the huggingface tokenizer. """ _tokenizer: PreTrainedTokenizerFast eos_token_id: int | None eos_token: str | None eos_token_ids: list[int] | set[int] | None bos_token_id: int | None bos_token: str | None vocab_size: int all_special_tokens: list[str] think_start: str | None think_end: str | None think_start_id: int | None think_end_id: int | None def __init__( self, tokenizer: Any, detokenizer_class: Any = ..., eos_token_ids: list[int] | set[int] | None = ..., chat_template: Any = ..., tool_parser: Any = ..., tool_call_start: str | None = ..., tool_call_end: str | None = ..., ) -> None: ... def encode(self, text: str, **kwargs: Any) -> list[int]: ... def decode(self, token_ids: list[int], **kwargs: Any) -> str: ... def apply_chat_template( self, messages: list[dict[str, Any]], tokenize: bool = False, add_generation_prompt: bool = False, tools: Any = None, **kwargs: Any, ) -> str: ... def get_vocab(self) -> dict[str, int]: ... def add_eos_token(self, token: str) -> None: ... @property def has_thinking(self) -> bool: ... @property def think_start(self) -> str | None: ... @property def think_end(self) -> str | None: ... @property def has_tool_calling(self) -> bool: ... @property def tool_call_start(self) -> str | None: ... @property def tool_call_end(self) -> str | None: ... @property def detokenizer(self) -> NaiveStreamingDetokenizer: """Get a stateful streaming detokenizer.""" def __getattr__(self, attr: str) -> Any: ... def __setattr__(self, attr: str, value: Any) -> None: ... class NewlineTokenizer(PreTrainedTokenizerFast): """A tokenizer that replaces newlines with and with new line.""" def __init__(self, *args, **kwargs) -> None: ... def encode(self, text, **kwargs): # -> list[int]: ... def encode_batch(self, texts, **kwargs): ... def decode(self, *args, **kwargs): # -> str: ... def batch_decode(self, *args, **kwargs): # -> list[str]: ... def load( model_path: Path, tokenizer_config_extra: dict[str, Any] | None = None, eos_token_ids: list[int] | int | None = None, ) -> TokenizerWrapper: """Load a huggingface tokenizer and try to infer the type of streaming detokenizer to use. Note, to use a fast streaming tokenizer, pass a local file path rather than a Hugging Face repo ID. """ # Alias for backward compatibility load_tokenizer = load def no_bos_or_eos(sequence: list[int], bos: int, eos: int) -> list[int]: ... ================================================ FILE: .mlx_typings/mlx_lm/utils.pyi ================================================ """ This type stub file was generated by pyright. """ import os from pathlib import Path from typing import Any, Callable, Dict, Optional, Tuple, Type, Union import mlx.nn as nn from transformers.utils.auto_docstring import ModelArgs from .tokenizer_utils import TokenizerWrapper if os.getenv("MLXLM_USE_MODELSCOPE", "False").lower() == "true": ... else: ... MODEL_REMAPPING = ... MAX_FILE_SIZE_GB = ... def compute_bits_per_weight(model): ... def hf_repo_to_path(hf_repo): # -> Path: ... def load_config(model_path: Path) -> dict: ... def load_model( model_path: Path, lazy: bool = False, strict: bool = True, model_config: dict[str, Any] = {}, get_model_classes: Callable[ [dict[str, Any]], Tuple[Type[nn.Module], Type[ModelArgs]] ] = ..., ) -> Tuple[nn.Module, dict[str, Any]]: """ Load and initialize the model from a given path. Args: model_path (Path): The path to load the model from. lazy (bool): If False eval the model parameters to make sure they are loaded in memory before returning, otherwise they will be loaded when needed. Default: ``False`` strict (bool): Whether or not to raise an exception if weights don't match. Default: ``True`` model_config (dict, optional): Optional configuration parameters for the model. Defaults to an empty dictionary. get_model_classes (Callable[[dict], Tuple[Type[nn.Module], Type]], optional): A function that returns the model class and model args class given a config. Defaults to the ``_get_classes`` function. Returns: Tuple[nn.Module, dict[str, Any]]: The loaded and initialized model and config. Raises: FileNotFoundError: If the weight files (.safetensors) are not found. ValueError: If the model class or args class are not found or cannot be instantiated. """ def load( path_or_hf_repo: str, tokenizer_config=..., model_config=..., adapter_path: Optional[str] = ..., lazy: bool = ..., return_config: bool = ..., revision: str = ..., ) -> Union[ Tuple[nn.Module, TokenizerWrapper], Tuple[nn.Module, TokenizerWrapper, Dict[str, Any]], ]: """ Load the model and tokenizer from a given path or a huggingface repository. Args: path_or_hf_repo (Path): The path or the huggingface repository to load the model from. tokenizer_config (dict, optional): Configuration parameters specifically for the tokenizer. Defaults to an empty dictionary. model_config(dict, optional): Configuration parameters specifically for the model. Defaults to an empty dictionary. adapter_path (str, optional): Path to the LoRA adapters. If provided, applies LoRA layers to the model. Default: ``None``. lazy (bool): If ``False`` eval the model parameters to make sure they are loaded in memory before returning, otherwise they will be loaded when needed. Default: ``False`` return_config (bool: If ``True`` return the model config as the last item.. revision (str, optional): A revision id which can be a branch name, a tag, or a commit hash. Returns: Union[Tuple[nn.Module, TokenizerWrapper], Tuple[nn.Module, TokenizerWrapper, Dict[str, Any]]]: A tuple containing the loaded model, tokenizer and, if requested, the model config. Raises: FileNotFoundError: If config file or safetensors are not found. ValueError: If model class or args class are not found. """ def make_shards(weights: dict, max_file_size_gb: int = ...) -> list: """ Splits the weights into smaller shards. Args: weights (dict): Model weights. max_file_size_gb (int): Maximum size of each shard in gigabytes. Returns: list: List of weight shards. """ def create_model_card( path: Union[str, Path], hf_path: Union[str, Path, None] ): # -> None: """ Uploads the model to Hugging Face hub. Args: path (Union[str, Path]): Local path to the model. hf_path (Union[str, Path, None]): Path to the original Hugging Face model. """ def upload_to_hub(path: str, upload_repo: str): # -> None: """ Uploads the model to Hugging Face hub. Args: path (str): Local path to the model. upload_repo (str): Name of the HF repo to upload to. """ def save_model( save_path: Union[str, Path], model: nn.Module, *, donate_model: bool = ... ) -> None: """Save model weights and metadata index into specified directory.""" def quantize_model( model: nn.Module, config: dict, group_size: int, bits: int, mode: str = ..., quant_predicate: Optional[Callable[[str, nn.Module], Union[bool, dict]]] = ..., ) -> Tuple[nn.Module, dict]: """ Applies quantization to the model weights. Args: model (nn.Module): The model to be quantized. config (dict): Model configuration. group_size (int): Group size for quantization. bits (int): Bits per weight for quantization. mode (str): The quantization mode. quant_predicate (Callable): A callable that decides how to quantize each layer based on the path. Accepts the layer `path` and the `module`. Returns either a bool to signify quantize/no quantize or a dict of quantization parameters to pass to `to_quantized`. Returns: Tuple: Tuple containing quantized model and config. """ def save_config(config: dict, config_path: Union[str, Path]) -> None: """Save the model configuration to the ``config_path``. The final configuration will be sorted before saving for better readability. Args: config (dict): The model configuration. config_path (Union[str, Path]): Model configuration file path. """ def save( dst_path: Union[str, Path], src_path_or_repo: Union[str, Path], model: nn.Module, tokenizer: TokenizerWrapper, config: Dict[str, Any], donate_model: bool = ..., ): # -> None: ... def common_prefix_len(list1, list2): # -> int: """ Calculates the length of the common prefix of two lists. Args: list1: The first list of strings. list2: The second list of strings. Returns: The length of the common prefix. Returns 0 if lists are empty or do not match at the first element. """ def does_model_support_input_embeddings(model: nn.Module) -> bool: """ Check if the model supports input_embeddings in its call signature. Args: model (nn.Module): The model to check. Returns: bool: True if the model supports input_embeddings, False otherwise. """ ================================================ FILE: .python-version ================================================ 3.13 ================================================ FILE: .swift-format ================================================ { "version": 1, "indentation": { "spaces": 4 } } ================================================ FILE: .vscode/extensions.json ================================================ { "recommendations": [ "detachhead.basedpyright", "ms-python.python" ], "unwantedRecommendations": [ "ms-python.vscode-pylance", "ms-python.pyright", "ms-python.mypy-type-checker" ] } ================================================ FILE: .vscode/settings.json ================================================ { "basedpyright.importStrategy": "fromEnvironment" } ================================================ FILE: .zed/settings.json ================================================ // Folder-specific settings // // For a full list of overridable settings, and general information on folder-specific settings, // see the documentation: https://zed.dev/docs/configuring-zed#settings-files { "lsp": { "nix_python": { "binary": { "path": "nix", "arguments": [ "run", "--quiet", "--no-warn-dirty", "--no-allow-import-from-derivation", "--print-build-logs", "never", "${projectRoot}#python-lsp", "--", "--stdio" ] } } }, "languages": { "Python": { "language_servers": ["nix_python"] } } } ================================================ FILE: AGENTS.md ================================================ # AGENTS.md This file provides guidance to AI coding agents when working with code in this repository. ## Project Overview exo is a distributed AI inference system that connects multiple devices into a cluster. It enables running large language models across multiple machines using MLX as the inference backend and libp2p for peer-to-peer networking. ## Build & Run Commands ```bash # Build the dashboard (required before running exo) cd dashboard && npm install && npm run build && cd .. # Run exo (starts both master and worker with API at http://localhost:52415) uv run exo # Run with verbose logging uv run exo -v # or -vv for more verbose # Run tests (excludes slow tests by default) uv run pytest # Run all tests including slow tests uv run pytest -m "" # Run a specific test file uv run pytest src/exo/shared/tests/test_election.py # Run a specific test function uv run pytest src/exo/shared/tests/test_election.py::test_function_name # Type checking (strict mode) uv run basedpyright # Linting uv run ruff check # Format code (using nix) nix fmt ``` ## Pre-Commit Checks (REQUIRED) **IMPORTANT: Always run these checks before committing code. CI will fail if these don't pass.** ```bash # 1. Type checking - MUST pass with 0 errors uv run basedpyright # 2. Linting - MUST pass uv run ruff check # 3. Formatting - MUST be applied nix fmt # 4. Tests - MUST pass uv run pytest ``` Run all checks in sequence: ```bash uv run basedpyright && uv run ruff check && nix fmt && uv run pytest ``` If `nix fmt` changes any files, stage them before committing. The CI runs `nix flake check` which verifies formatting, linting, and runs Rust tests. ## Architecture ### Node Composition A single exo `Node` (src/exo/main.py) runs multiple components: - **Router**: libp2p-based pub/sub messaging via Rust bindings (exo_pyo3_bindings) - **Worker**: Handles inference tasks, downloads models, manages runner processes - **Master**: Coordinates cluster state, places model instances across nodes - **Election**: Bully algorithm for master election - **API**: FastAPI server for OpenAI-compatible chat completions ### Message Flow Components communicate via typed pub/sub topics (src/exo/routing/topics.py): - `GLOBAL_EVENTS`: Master broadcasts indexed events to all workers - `LOCAL_EVENTS`: Workers send events to master for indexing - `COMMANDS`: Workers/API send commands to master - `ELECTION_MESSAGES`: Election protocol messages - `CONNECTION_MESSAGES`: libp2p connection updates ### Event Sourcing The system uses event sourcing for state management: - `State` (src/exo/shared/types/state.py): Immutable state object - `apply()` (src/exo/shared/apply.py): Pure function that applies events to state - Master indexes events and broadcasts; workers apply indexed events ### Key Type Hierarchy - `src/exo/shared/types/`: Pydantic models for all shared types - `events.py`: Event types (discriminated union) - `commands.py`: Command types - `tasks.py`: Task types for worker execution - `state.py`: Cluster state model ### Rust Components Rust code in `rust/` provides: - `networking`: libp2p networking (gossipsub, peer discovery) - `exo_pyo3_bindings`: PyO3 bindings exposing Rust to Python - `system_custodian`: System-level operations ### Dashboard Svelte 5 + TypeScript frontend in `dashboard/`. Build output goes to `dashboard/build/` and is served by the API. ## Code Style Requirements From .cursorrules: - Strict, exhaustive typing - never bypass the type-checker - Use `Literal[...]` for enum-like sets, `typing.NewType` for primitives - Pydantic models with `frozen=True` and `strict=True` - Pure functions with injectable effect handlers for side-effects - Descriptive names - no abbreviations or 3-letter acronyms - Catch exceptions only where you can handle them meaningfully - Use `@final` and immutability wherever applicable ## Testing Tests use pytest-asyncio with `asyncio_mode = "auto"`. Tests are in `tests/` subdirectories alongside the code they test. The `EXO_TESTS=1` env var is set during tests. ## Dashboard UI Testing & Screenshots ### Building and Running the Dashboard ```bash # Build the dashboard (must be done before running exo) cd dashboard && npm install && npm run build && cd .. # Start exo (serves the dashboard at http://localhost:52415) uv run exo & sleep 8 # Wait for server to start ``` ### Taking Headless Screenshots with Playwright Use Playwright with headless Chromium for programmatic screenshots — no manual browser interaction needed. **Setup (one-time):** ```bash npx --yes playwright install chromium cd /tmp && npm init -y && npm install playwright ``` **Taking screenshots:** ```javascript // Run from /tmp where playwright is installed: cd /tmp && node -e "..." const { chromium } = require('playwright'); (async () => { const browser = await chromium.launch({ headless: true }); const page = await browser.newPage({ viewport: { width: 1280, height: 800 } }); await page.goto('http://localhost:52415', { waitUntil: 'networkidle' }); await page.waitForTimeout(2000); // Inject test data into localStorage if needed (e.g., recent models) await page.evaluate(() => { localStorage.setItem('exo-recent-models', JSON.stringify([ { modelId: 'mlx-community/Qwen3-30B-A3B-4bit', launchedAt: Date.now() }, ])); }); await page.reload({ waitUntil: 'networkidle' }); await page.waitForTimeout(2000); // Interact with UI elements await page.locator('text=SELECT MODEL').click(); await page.waitForTimeout(1000); // Take screenshot await page.screenshot({ path: '/tmp/screenshot.png', fullPage: false }); await browser.close(); })(); ``` ### Uploading Images to GitHub PRs GitHub's API doesn't support direct image upload for PR comments. Workaround: 1. **Commit images to the branch** (temporarily): ```bash cp /tmp/screenshot.png . git add screenshot.png git commit -m "temp: add screenshots for PR" git push origin COMMIT_SHA=$(git rev-parse HEAD) ``` 2. **Post PR comment** referencing the raw image URL (uses permanent commit SHA so images survive deletion): ```bash gh pr comment --body "![Screenshot](https://raw.githubusercontent.com/exo-explore/exo/${COMMIT_SHA}/screenshot.png)" ``` 3. **Remove the images** from the branch: ```bash git rm screenshot.png git commit -m "chore: remove temporary screenshot files" git push origin ``` The images still render in the PR comment because they reference the permanent commit SHA. ================================================ FILE: CONTRIBUTING.md ================================================ # Contributing to EXO Thank you for your interest in contributing to EXO! ## Getting Started To run EXO from source: **Prerequisites:** - [uv](https://github.com/astral-sh/uv) (for Python dependency management) ```bash brew install uv ``` - [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon) ```bash brew install macmon ``` ```bash git clone https://github.com/exo-explore/exo.git cd exo/dashboard npm install && npm run build && cd .. uv run exo ``` ## Development EXO is built with a mix of Rust, Python, and TypeScript (Svelte for the dashboard), and the codebase is actively evolving. Before starting work: - Pull the latest source to ensure you're working with the most recent code - Keep your changes focused - implement one feature or fix per pull request - Avoid combining unrelated changes, even if they seem small This makes reviews faster and helps us maintain code quality as the project evolves. ## Code Style Write pure functions where possible. When adding new code, prefer Rust unless there's a good reason otherwise. Leverage the type systems available to you - Rust's type system, Python type hints, and TypeScript types. Comments should explain why you're doing something, not what the code does - especially for non-obvious decisions. Run `nix fmt` to auto-format your code before submitting. ## Model Cards EXO uses TOML-based model cards to define model metadata and capabilities. Model cards are stored in: - `resources/inference_model_cards/` for text generation models - `resources/image_model_cards/` for image generation models - `~/.exo/custom_model_cards/` for user-added custom models ### Adding a Model Card To add a new model, create a TOML file with the following structure: ```toml model_id = "mlx-community/Llama-3.2-1B-Instruct-4bit" n_layers = 16 hidden_size = 2048 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "4bit" base_model = "Llama 3.2 1B" capabilities = ["text"] [storage_size] in_bytes = 729808896 ``` ### Required Fields - `model_id`: Hugging Face model identifier - `n_layers`: Number of transformer layers - `hidden_size`: Hidden dimension size - `supports_tensor`: Whether the model supports tensor parallelism - `tasks`: List of supported tasks (`TextGeneration`, `TextToImage`, `ImageToImage`) - `family`: Model family (e.g., "llama", "deepseek", "qwen") - `quantization`: Quantization level (e.g., "4bit", "8bit", "bf16") - `base_model`: Human-readable base model name - `capabilities`: List of capabilities (e.g., `["text"]`, `["text", "thinking"]`) ### Optional Fields - `components`: For multi-component models (like image models with separate text encoders and transformers) - `uses_cfg`: Whether the model uses classifier-free guidance (for image models) - `trust_remote_code`: Whether to allow remote code execution (defaults to `false` for security) ### Capabilities The `capabilities` field defines what the model can do: - `text`: Standard text generation - `thinking`: Model supports chain-of-thought reasoning - `thinking_toggle`: Thinking can be enabled/disabled via `enable_thinking` parameter - `image_edit`: Model supports image-to-image editing (FLUX.1-Kontext) ### Security Note By default, `trust_remote_code` is set to `false` for security. Only enable it if the model explicitly requires remote code execution from the Hugging Face hub. ## API Adapters EXO supports multiple API formats through an adapter pattern. Adapters convert API-specific request formats to the internal `TextGenerationTaskParams` format and convert internal token chunks back to API-specific responses. ### Adapter Architecture All adapters live in `src/exo/master/adapters/` and follow the same pattern: 1. Convert API-specific requests to `TextGenerationTaskParams` 2. Handle both streaming and non-streaming response generation 3. Convert internal `TokenChunk` objects to API-specific formats 4. Manage error handling and edge cases ### Existing Adapters - `chat_completions.py`: OpenAI Chat Completions API - `claude.py`: Anthropic Claude Messages API - `responses.py`: OpenAI Responses API - `ollama.py`: Ollama API (for OpenWebUI compatibility) ### Adding a New API Adapter To add support for a new API format: 1. Create a new adapter file in `src/exo/master/adapters/` 2. Implement a request conversion function: ```python def your_api_request_to_text_generation( request: YourAPIRequest, ) -> TextGenerationTaskParams: # Convert API request to internal format pass ``` 3. Implement streaming response generation: ```python async def generate_your_api_stream( command_id: CommandId, chunk_stream: AsyncGenerator[TokenChunk | ErrorChunk | ToolCallChunk, None], ) -> AsyncGenerator[str, None]: # Convert internal chunks to API-specific streaming format pass ``` 4. Implement non-streaming response collection: ```python async def collect_your_api_response( command_id: CommandId, chunk_stream: AsyncGenerator[TokenChunk | ErrorChunk | ToolCallChunk, None], ) -> AsyncGenerator[str]: # Collect all chunks and return single response pass ``` 5. Register the adapter endpoints in `src/exo/master/api.py` The adapter pattern keeps API-specific logic isolated from core inference systems. Internal systems (worker, runner, event sourcing) only see `TextGenerationTaskParams` and `TokenChunk` objects - no API-specific types cross the adapter boundary. For detailed API documentation, see [docs/api.md](docs/api.md). ## Testing EXO relies heavily on manual testing at this point in the project, but this is evolving. Before submitting a change, test both before and after to demonstrate how your change improves behavior. Do the best you can with the hardware you have available - if you need help testing, ask and we'll do our best to assist. Add automated tests where possible - we're actively working to substantially improve our automated testing story. ## Submitting Changes 1. Fork the repository 2. Create a feature branch (`git checkout -b feature/your-feature`) 3. Commit your changes (`git commit -am 'Add some feature'`) 4. Push to the branch (`git push origin feature/your-feature`) 5. Open a Pull Request and follow the PR template ## Reporting Issues If you find a bug or have a feature request, please open an issue on GitHub with: - A clear description of the problem or feature - Steps to reproduce (for bugs) - Expected vs actual behavior - Your environment (macOS version, hardware, etc.) ## Questions? Join our community: - [X](https://x.com/exolabs) ================================================ FILE: Cargo.toml ================================================ [workspace] resolver = "3" members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"] [workspace.package] version = "0.0.1" edition = "2024" [profile.dev] opt-level = 1 debug = true [profile.release] opt-level = 3 # Common shared dependendencies configured once at the workspace # level, to be re-used more easily across workspace member crates. # # Common configurations include versions, paths, features, etc. [workspace.dependencies] ## Crate members as common dependencies networking = { path = "rust/networking" } util = { path = "rust/util" } # Macro dependecies extend = "1.2" delegate = "0.13" # Utility dependencies keccak-const = "0.2" # Async dependencies async-stream = "0.3" tokio = "1.46" futures-lite = "2.6.1" futures-timer = "3.0" # Data structures either = "1.15" # Tracing/logging log = "0.4" # networking libp2p = "0.56" libp2p-tcp = "0.44" [workspace.lints.rust] static_mut_refs = "warn" # Or use "warn" instead of deny incomplete_features = "allow" # Clippy's lint category level configurations; # every member crate needs to inherit these by adding # # ```toml # [lints] # workspace = true # ``` # # to their `Cargo.toml` files [workspace.lints.clippy] # Clippy lint categories meant to be enabled all at once correctness = { level = "deny", priority = -1 } suspicious = { level = "warn", priority = -1 } style = { level = "warn", priority = -1 } complexity = { level = "warn", priority = -1 } perf = { level = "warn", priority = -1 } pedantic = { level = "warn", priority = -1 } nursery = { level = "warn", priority = -1 } cargo = { level = "warn", priority = -1 } # Individual Clippy lints from the `restriction` category arithmetic_side_effects = "warn" as_conversions = "warn" assertions_on_result_states = "warn" clone_on_ref_ptr = "warn" decimal_literal_representation = "warn" default_union_representation = "warn" deref_by_slicing = "warn" disallowed_script_idents = "deny" else_if_without_else = "warn" empty_enum_variants_with_brackets = "warn" empty_structs_with_brackets = "warn" error_impl_error = "warn" exit = "deny" expect_used = "warn" float_cmp_const = "warn" get_unwrap = "warn" if_then_some_else_none = "warn" impl_trait_in_params = "warn" indexing_slicing = "warn" infinite_loop = "warn" let_underscore_must_use = "warn" let_underscore_untyped = "warn" lossy_float_literal = "warn" mem_forget = "warn" missing_inline_in_public_items = "warn" multiple_inherent_impl = "warn" multiple_unsafe_ops_per_block = "warn" mutex_atomic = "warn" non_zero_suggestions = "warn" panic = "warn" partial_pub_fields = "warn" pattern_type_mismatch = "warn" pub_without_shorthand = "warn" rc_buffer = "warn" rc_mutex = "warn" redundant_type_annotations = "warn" renamed_function_params = "warn" rest_pat_in_fully_bound_structs = "warn" same_name_method = "warn" self_named_module_files = "deny" semicolon_inside_block = "warn" shadow_same = "warn" shadow_unrelated = "warn" str_to_string = "warn" string_add = "warn" string_lit_chars_any = "warn" string_to_string = "warn" tests_outside_test_module = "warn" todo = "warn" try_err = "warn" undocumented_unsafe_blocks = "warn" unnecessary_safety_comment = "warn" unnecessary_safety_doc = "warn" unneeded_field_pattern = "warn" unseparated_literal_suffix = "warn" unused_result_ok = "warn" unused_trait_names = "warn" unwrap_used = "warn" verbose_file_reads = "warn" ================================================ FILE: LICENSE ================================================ Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "[]" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright 2025 Exo Technologies Ltd Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ================================================ FILE: MISSED_THINGS.md ================================================ # Missed things [X] Log EXO_LIBP2P_NAMESPACE on start in exo/main.py [X] Ordering of warmup was changed, which is wrong. It was changed to rank < n-1, then rank=n-1. It should be rank!=0 then rank=0 (this matches the auto_parallel implementation. NOTE: we use a different convention to mlx-lm, our terminal rank is rank=n-1 whereas mlx-lm is rank=0 hence i can see why this was changed wrongly). [X] Downloads keying by model_id not shard_metadata (worker/plan.py, worker/main.py). [X] Fetching download status of all models on start [X] Deduplication of tasks in plan_step. [X] resolve_allow_patterns should just be wildcard now. [X] no mx_barrier in genreate.py mlx_generate at the end. [] cache assertion not needed in auto_parallel.py PipelineLastLayer. [X] GPTOSS support dropped in auto_parallel.py. [X] sharding changed "all-to-sharded" became _all_to_sharded in auto_parallel.py. [X] same as above with "sharded-to-all" became _sharded_to_all in auto_parallel.py. [X] Dropped support for Ministral3Model, DeepseekV32Model, Glm4MoeModel, Qwen3NextModel, GptOssMode in auto_parallel.py. [] Dropped prefill/decode code in auto_parallel.py and utils_mlx.py. [X] KV_CACHE_BITS should be None to disable quantized KV cache. [X] Dropped _set_nofile_limit in utils_mlx.py. [X] We have group optional in load_mlx_items in utils_mlx.py. [X] Dropped add_missing_chat_templates for GptOss in load_mlx_items in utils_mlx.py. [X] Dropped model.make_cache in make_kv_cache in utils_mlx.py. [X] We put cache limit back in utils_mlx.py. [X] topology.py remove_node removes the connections after checking if node is is in self._node_id_to_rx_id_map. on beta_1 it checks after, so would remove stale connections I guess? [X] Missing Glm 4.7 model cards (this isn't ready yet but should be picked up, probably create an issue... the blocker is transforemrs version doesn't support the tokenizer for Glm 4.7. rc-1 does but we can't upgrade as it breaks other things.) [] try-except in _command_processor only excepts ValueError. This was silently failing leading to un-debuggable errors (we had a KeyError that was happening ). Changed this to catch Exception instead of ValueError. See exo-v2 89ae38405e0052e3c22405daf094b065878aa873 and fb99fea69b5a39017efc90c5dad0072e677455f0. [X] In placement.py, place_instance no longer looks at model_meta.supports_tensor and check if this tensor parallel number of nodes is supported by the model's tensor dimensions. [X] In placement.py, place_instanec, we no longer have the special case to exclude DeepSeek v3.1 pipeline parallel (it doesn't work). [] logger.warning("You have likely selected ibv for a single node instance; falling back to MlxRing") was changed to debug. That will spam this warning since it happens every time we query instance previews. [X] In placement_utils.py, get_mlx_jaccl_coordinators, We no longer prioritise Jaccl Coordinator IP. Now it picks the first one, which is unstable (Jaccl coordinator over TB5 is unstable). [X] Downloads keying by model_id not shard_metadata (worker/plan.py, worker/main.py). [X] Fetching download status of all models on start [X] Deduplication of tasks in plan_step. [X] resolve_allow_patterns should just be wildcard now. [X] KV_CACHE_BITS should be None to disable quantized KV cache. [X] We put cache limit back in utils_mlx.py. [X] In placement.py, place_instance no longer looks at model_meta.supports_tensor and check if this tensor parallel number of nodes is supported by the model's tensor dimensions. [X] In placement.py, place_instanec, we no longer have the special case to exclude DeepSeek v3.1 pipeline parallel (it doesn't work). [X] In placement_utils.py, get_mlx_jaccl_coordinators, We no longer prioritise Jaccl Coordinator IP. Now it picks the first one, which is unstable (Jaccl coordinator over TB5 is unstable). ================================================ FILE: PLATFORMS.md ================================================ # EXO Platform support (partial roadmap) ## Tier 1 support - tested and maintained Apple Silicon MacOS - Mac Studio: M3 Ultra - Mac Mini: M4 Pro - Macbook Pro: M5, M4 Max ## Tier 2 support - checked occasionally, should run without crashing ## Tier 3 support - minimal support and testing, but no theoretical reason it shouldnt work # Planned ## Tier 1 Linux CUDA Support - Nvidia DGX Spark Linux CPU Support ## Tier 2 Linux Vulkan Support -- depends heavily on ecosystem - Framework Desktop Linux CUDA Support -- depends heavily on ecosystem - Framework Desktop ## Longer term! Windows CUDA Support Windows CPU Support ================================================ FILE: README.md ================================================
exo logo exo: Run frontier AI locally. Maintained by [exo labs](https://x.com/exolabs).

Discord X License: Apache-2.0

--- exo connects all your devices into an AI cluster. Not only does exo enable running models larger than would fit on a single device, but with [day-0 support for RDMA over Thunderbolt](https://x.com/exolabs/status/2001817749744476256?s=20), makes models run faster as you add more devices. ## Features - **Automatic Device Discovery**: Devices running exo automatically discover each other - no manual configuration. - **RDMA over Thunderbolt**: exo ships with [day-0 support for RDMA over Thunderbolt 5](https://x.com/exolabs/status/2001817749744476256?s=20), enabling 99% reduction in latency between devices. - **Topology-Aware Auto Parallel**: exo figures out the best way to split your model across all available devices based on a realtime view of your device topology. It takes into account device resources and network latency/bandwidth between each link. - **Tensor Parallelism**: exo supports sharding models, for up to 1.8x speedup on 2 devices and 3.2x speedup on 4 devices. - **MLX Support**: exo uses [MLX](https://github.com/ml-explore/mlx) as an inference backend and [MLX distributed](https://ml-explore.github.io/mlx/build/html/usage/distributed.html) for distributed communication. - **Multiple API Compatibility**: Compatible with OpenAI Chat Completions API, Claude Messages API, OpenAI Responses API, and Ollama API - use your existing tools and clients. - **Custom Model Support**: Load custom models from HuggingFace hub to expand the range of available models. ## Dashboard exo includes a built-in dashboard for managing your cluster and chatting with models.

exo dashboard - cluster view showing 4 x M3 Ultra Mac Studio with DeepSeek v3.1 and Kimi-K2-Thinking loaded

4 × 512GB M3 Ultra Mac Studio running DeepSeek v3.1 (8-bit) and Kimi-K2-Thinking (4-bit)

## Benchmarks
Qwen3-235B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA Benchmark - Qwen3-235B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA

Source: Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5

DeepSeek v3.1 671B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA Benchmark - DeepSeek v3.1 671B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA

Source: Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5

Kimi K2 Thinking (native 4-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA Benchmark - Kimi K2 Thinking (native 4-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA

Source: Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5

--- ## Quick Start Devices running exo automatically discover each other, without needing any manual configuration. Each device provides an API and a dashboard for interacting with your cluster (runs at `http://localhost:52415`). There are two ways to run exo: ### Run from Source (macOS) If you have [Nix](https://nixos.org/) installed, you can skip most of the steps below and run exo directly: ```bash nix run .#exo ``` **Note:** To accept the Cachix binary cache (and avoid the Xcode Metal ToolChain), add to `/etc/nix/nix.conf`: ``` trusted-users = root (or your username) experimental-features = nix-command flakes ``` Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-daemon` **Prerequisites:** - [Xcode](https://developer.apple.com/xcode/) (provides the Metal ToolChain required for MLX compilation) - [brew](https://github.com/Homebrew/brew) (for simple package management on macOS) ```bash /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)" ``` - [uv](https://github.com/astral-sh/uv) (for Python dependency management) - [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon) - [node](https://github.com/nodejs/node) (for building the dashboard) ```bash brew install uv macmon node ``` - [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now) ```bash curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh rustup toolchain install nightly ``` Clone the repo, build the dashboard, and run exo: ```bash # Clone exo git clone https://github.com/exo-explore/exo # Build dashboard cd exo/dashboard && npm install && npm run build && cd .. # Run exo uv run exo ``` This starts the exo dashboard and API at http://localhost:52415/ *Please view the section on RDMA to enable this feature on MacOS >=26.2!* ### Run from Source (Linux) **Prerequisites:** - [uv](https://github.com/astral-sh/uv) (for Python dependency management) - [node](https://github.com/nodejs/node) (for building the dashboard) - version 18 or higher - [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now) **Installation methods:** **Option 1: Using system package manager (Ubuntu/Debian example):** ```bash # Install Node.js and npm sudo apt update sudo apt install nodejs npm # Install uv curl -LsSf https://astral.sh/uv/install.sh | sh # Install Rust (using rustup) curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh rustup toolchain install nightly ``` **Option 2: Using Homebrew on Linux (if preferred):** ```bash # Install Homebrew on Linux /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)" # Install dependencies brew install uv node # Install Rust (using rustup) curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh rustup toolchain install nightly ``` **Note:** The `macmon` package is macOS-only and not required for Linux. Clone the repo, build the dashboard, and run exo: ```bash # Clone exo git clone https://github.com/exo-explore/exo # Build dashboard cd exo/dashboard && npm install && npm run build && cd .. # Run exo uv run exo ``` This starts the exo dashboard and API at http://localhost:52415/ **Important note for Linux users:** Currently, exo runs on CPU on Linux. GPU support for Linux platforms is under development. If you'd like to see support for your specific Linux hardware, please [search for existing feature requests](https://github.com/exo-explore/exo/issues) or create a new one. **Configuration Options:** - `--no-worker`: Run exo without the worker component. Useful for coordinator-only nodes that handle networking and orchestration but don't execute inference tasks. This is helpful for machines without sufficient GPU resources but with good network connectivity. ```bash uv run exo --no-worker ``` **File Locations (Linux):** exo follows the [XDG Base Directory Specification](https://specifications.freedesktop.org/basedir-spec/basedir-spec-latest.html) on Linux: - **Configuration files**: `~/.config/exo/` (or `$XDG_CONFIG_HOME/exo/`) - **Data files**: `~/.local/share/exo/` (or `$XDG_DATA_HOME/exo/`) - **Cache files**: `~/.cache/exo/` (or `$XDG_CACHE_HOME/exo/`) - **Log files**: `~/.cache/exo/exo_log/` (with automatic log rotation) - **Custom model cards**: `~/.local/share/exo/custom_model_cards/` You can override these locations by setting the corresponding XDG environment variables. ### macOS App exo ships a macOS app that runs in the background on your Mac. exo macOS App - running on a MacBook The macOS app requires macOS Tahoe 26.2 or later. Download the latest build here: [EXO-latest.dmg](https://assets.exolabs.net/EXO-latest.dmg). The app will ask for permission to modify system settings and install a new Network profile. Improvements to this are being worked on. **Custom Namespace for Cluster Isolation:** The macOS app includes a custom namespace feature that allows you to isolate your exo cluster from others on the same network. This is configured through the `EXO_LIBP2P_NAMESPACE` setting: - **Use cases**: - Running multiple separate exo clusters on the same network - Isolating development/testing clusters from production clusters - Preventing accidental cluster joining - **Configuration**: Access this setting in the app's Advanced settings (or set the `EXO_LIBP2P_NAMESPACE` environment variable when running from source) The namespace is logged on startup for debugging purposes. #### Uninstalling the macOS App The recommended way to uninstall is through the app itself: click the menu bar icon → Advanced → Uninstall. This cleanly removes all system components. If you've already deleted the app, you can run the standalone uninstaller script: ```bash sudo ./app/EXO/uninstall-exo.sh ``` This removes: - Network setup LaunchDaemon - Network configuration script - Log files - The "exo" network location **Note:** You'll need to manually remove EXO from Login Items in System Settings → General → Login Items. --- ### Enabling RDMA on macOS RDMA is a new capability added to macOS 26.2. It works on any Mac with Thunderbolt 5 (M4 Pro Mac Mini, M4 Max Mac Studio, M4 Max MacBook Pro, M3 Ultra Mac Studio). Please refer to the caveats for immediate troubleshooting. To enable RDMA on macOS, follow these steps: 1. Shut down your Mac. 2. Hold down the power button for 10 seconds until the boot menu appears. 3. Select "Options" to enter Recovery mode. 4. When the Recovery UI appears, open the Terminal from the Utilities menu. 5. In the Terminal, type: ``` rdma_ctl enable ``` and press Enter. 6. Reboot your Mac. After that, RDMA will be enabled in macOS and exo will take care of the rest. **Important Caveats** 1. Devices that wish to be part of an RDMA cluster must be connected to all other devices in the cluster. 2. The cables must support TB5. 3. On a Mac Studio, you cannot use the Thunderbolt 5 port next to the Ethernet port. 4. If running from source, please use the script found at `tmp/set_rdma_network_config.sh`, which will disable Thunderbolt Bridge and set dhcp on each RDMA port. 5. RDMA ports may be unable to discover each other on different versions of MacOS. Please ensure that OS versions match exactly (even beta version numbers) on all devices. --- ## Environment Variables exo supports several environment variables for configuration: | Variable | Description | Default | |----------|-------------|---------| | `EXO_MODELS_PATH` | Colon-separated paths to search for pre-downloaded models (e.g., on NFS mounts or shared storage) | None | | `EXO_MODELS_DIR` | Directory where exo downloads and stores models | `~/.local/share/exo/models` (Linux) or `~/.exo/models` (macOS) | | `EXO_OFFLINE` | Run without internet connection (uses only local models) | `false` | | `EXO_ENABLE_IMAGE_MODELS` | Enable image model support | `false` | | `EXO_LIBP2P_NAMESPACE` | Custom namespace for cluster isolation | None | | `EXO_FAST_SYNCH` | Control MLX_METAL_FAST_SYNCH behavior (for JACCL backend) | Auto | | `EXO_TRACING_ENABLED` | Enable distributed tracing for performance analysis | `false` | **Example usage:** ```bash # Use pre-downloaded models from NFS mount EXO_MODELS_PATH=/mnt/nfs/models:/opt/ai-models uv run exo # Run in offline mode EXO_OFFLINE=true uv run exo # Enable image models EXO_ENABLE_IMAGE_MODELS=true uv run exo # Use custom namespace for cluster isolation EXO_LIBP2P_NAMESPACE=my-dev-cluster uv run exo ``` --- ### Using the API exo provides multiple API-compatible interfaces for maximum compatibility with existing tools: - **OpenAI Chat Completions API** - Compatible with OpenAI clients - **Claude Messages API** - Compatible with Anthropic's Claude format - **OpenAI Responses API** - Compatible with OpenAI's Responses format - **Ollama API** - Compatible with Ollama and tools like OpenWebUI If you prefer to interact with exo via the API, here is an example creating an instance of a small model (`mlx-community/Llama-3.2-1B-Instruct-4bit`), sending a chat completions request and deleting the instance. --- **1. Preview instance placements** The `/instance/previews` endpoint will preview all valid placements for your model. ```bash curl "http://localhost:52415/instance/previews?model_id=llama-3.2-1b" ``` Sample response: ```json { "previews": [ { "model_id": "mlx-community/Llama-3.2-1B-Instruct-4bit", "sharding": "Pipeline", "instance_meta": "MlxRing", "instance": {...}, "memory_delta_by_node": {"local": 729808896}, "error": null } // ...possibly more placements... ] } ``` This will return all valid placements for this model. Pick a placement that you like. To pick the first one, pipe into `jq`: ```bash curl "http://localhost:52415/instance/previews?model_id=llama-3.2-1b" | jq -c '.previews[] | select(.error == null) | .instance' | head -n1 ``` --- **2. Create a model instance** Send a POST to `/instance` with your desired placement in the `instance` field (the full payload must match types as in `CreateInstanceParams`), which you can copy from step 1: ```bash curl -X POST http://localhost:52415/instance \ -H 'Content-Type: application/json' \ -d '{ "instance": {...} }' ``` Sample response: ```json { "message": "Command received.", "command_id": "e9d1a8ab-...." } ``` --- **3. Send a chat completion** Now, make a POST to `/v1/chat/completions` (the same format as OpenAI's API): ```bash curl -N -X POST http://localhost:52415/v1/chat/completions \ -H 'Content-Type: application/json' \ -d '{ "model": "mlx-community/Llama-3.2-1B-Instruct-4bit", "messages": [ {"role": "user", "content": "What is Llama 3.2 1B?"} ], "stream": true }' ``` --- **4. Delete the instance** When you're done, delete the instance by its ID (find it via `/state` or `/instance` endpoints): ```bash curl -X DELETE http://localhost:52415/instance/YOUR_INSTANCE_ID ``` ### Claude Messages API Compatibility Use the Claude Messages API format with the `/v1/messages` endpoint: ```bash curl -N -X POST http://localhost:52415/v1/messages \ -H 'Content-Type: application/json' \ -d '{ "model": "mlx-community/Llama-3.2-1B-Instruct-4bit", "messages": [ {"role": "user", "content": "Hello"} ], "max_tokens": 1024, "stream": true }' ``` ### OpenAI Responses API Compatibility Use the OpenAI Responses API format with the `/v1/responses` endpoint: ```bash curl -N -X POST http://localhost:52415/v1/responses \ -H 'Content-Type: application/json' \ -d '{ "model": "mlx-community/Llama-3.2-1B-Instruct-4bit", "messages": [ {"role": "user", "content": "Hello"} ], "stream": true }' ``` ### Ollama API Compatibility exo supports Ollama API endpoints for compatibility with tools like OpenWebUI: ```bash # Ollama chat curl -X POST http://localhost:52415/ollama/api/chat \ -H 'Content-Type: application/json' \ -d '{ "model": "mlx-community/Llama-3.2-1B-Instruct-4bit", "messages": [ {"role": "user", "content": "Hello"} ], "stream": false }' # List models (Ollama format) curl http://localhost:52415/ollama/api/tags ``` ### Custom Model Loading from HuggingFace You can add custom models from the HuggingFace hub: ```bash curl -X POST http://localhost:52415/models/add \ -H 'Content-Type: application/json' \ -d '{ "model_id": "mlx-community/my-custom-model" }' ``` **Security Note:** Custom models requiring `trust_remote_code` in their configuration must be explicitly enabled (default is false) for security. Only enable this if you trust the model's remote code execution. Models are fetched from HuggingFace and stored locally as custom model cards. **Other useful API endpoints*:** - List all models: `curl http://localhost:52415/models` - List downloaded models only: `curl http://localhost:52415/models?status=downloaded` - Search HuggingFace: `curl "http://localhost:52415/models/search?query=llama&limit=10"` - Inspect instance IDs and deployment state: `curl http://localhost:52415/state` For further details, see: - API documentation in [docs/api.md](docs/api.md). - API types and endpoints in [src/exo/master/api.py](src/exo/master/api.py). --- ## Benchmarking The `exo-bench` tool measures model prefill and token generation speed across different placement configurations. This helps you optimize model performance and validate improvements. **Prerequisites:** - Nodes should be running with `uv run exo` before benchmarking - The tool uses the `/bench/chat/completions` endpoint **Basic usage:** ```bash uv run bench/exo_bench.py \ --model Llama-3.2-1B-Instruct-4bit \ --pp 128,256,512 \ --tg 128,256 ``` **Key parameters:** - `--model`: Model to benchmark (short ID or HuggingFace ID) - `--pp`: Prompt size hints (comma-separated integers) - `--tg`: Generation lengths (comma-separated integers) - `--max-nodes`: Limit placements to N nodes (default: 4) - `--instance-meta`: Filter by `ring`, `jaccl`, or `both` (default: both) - `--sharding`: Filter by `pipeline`, `tensor`, or `both` (default: both) - `--repeat`: Number of repetitions per configuration (default: 1) - `--warmup`: Warmup runs per placement (default: 0) - `--json-out`: Output file for results (default: bench/results.json) **Example with filters:** ```bash uv run bench/exo_bench.py \ --model Llama-3.2-1B-Instruct-4bit \ --pp 128,512 \ --tg 128 \ --max-nodes 2 \ --sharding tensor \ --repeat 3 \ --json-out my-results.json ``` The tool outputs performance metrics including prompt tokens per second (prompt_tps), generation tokens per second (generation_tps), and peak memory usage for each configuration. --- ## Hardware Accelerator Support On macOS, exo uses the GPU. On Linux, exo currently runs on CPU. We are working on extending hardware accelerator support. If you'd like support for a new hardware platform, please [search for an existing feature request](https://github.com/exo-explore/exo/issues) and add a thumbs up so we know what hardware is important to the community. --- ## Contributing See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo. ================================================ FILE: RULES.md ================================================ # Repository Rules * if you see any code that violates these rules, raise it with me directly rather than trying to fix. * where applicable, file a GitHub Issue. * adhere to these rules strictly. ## General Rules * if its possible to eliminate an extra try-catch or if-statement at runtime using type-level discipline, do it! * name your types, functions, and classes appropriately. * no three-letter acronyms. * no non-standard contractions. * each data type has a meaning, pick a name which is accurate and descriptive. * the average layman should be able to easily understand what your function does using the function signature alone! * sometimes, there will be exceptions. eg, when you're using specific technical terms that are well understood (saga, event, etc). * usually, you'll think that your code is an exception to the rules, but it won't be. ## State, Functions and Classes * every function, given the same inputs, should produce the same outputs. ie, no hidden state. * use classes to prevent fixed state from being mutated arbitrarily (unsafely); methods provide a safe way of interfacing with state. * if your logic doesn't mutate fixed state, it probably belongs in a standalone function rather than a class. * functions shouldn't usually produce side-effects (they should be computationally pure). * if, for example, you're updating a state using an event (computationally pure), and you want to trigger a saga (computational side-effect), store the logic for triggering the saga into an effect handler (a function, capable of producing side-effects, that you pass into an otherwise computationally pure function, so that it may trigger side-effects safely). ## Pydantic * read the Pydantic docs. * respect the Pydantic docs. * pydantic is all you need. * declare and re-use a central `ConfigDict` for your use-case, you'll usually want `frozen` and `strict` to be `True`. ## Unique ID (UUID) Generation * inherit from Pydantic's `UUID4` class to create your own UUID class. * use `uuid.uuid4()` to initialize your class with a fresh UUID where possible. * ensure that idempotency tags are generated by taking the salted hash of persisted state. * rationale: if a node crashes and resumes from an older state, it should not accidentally re-publish the same event twice under different idempotency tags. * every distinct function should feature a unique salt, so that there are no accidental collisions in idempotency tags. ## Type Wrappers * reuse types that already exist in the Python standard library. * when two distinct data types are structurally identical (for example, different IDs which are both UUIDs but shouldn't never mixed up), make sure they can't be conflated by the type system. * if you're working with a primitive data type (`str`, `int`, etc), use `NewType` (it has zero runtime overhead). * if you're working with serializable data objects, consider adding a field (type `str`) that states its type. ## Type Discipline * do not bypass the type-checker, preserve strict typing by any means necessary. * by default, use literal types (like `Literal['one', 'two']`) where an enum seems appropriate. pro-tip: Python's type system is quite complex and feature-rich, so reading the documentation is often advisable; Matt discovered that Python `typing` library allows you to check that you've implemented a `match` exhaustively using `Literal` and `get_args(type)` after reading the docs. ## Use of `@final`, Freezing * use wherever applicable. ## Error Handling * don't try-catch for no reason. * make sure that you always know where and when the exceptions your code produces are meant to be handled, so that it's never a nasty surprise. * always write the rationale for your error-handling down in the docstring! * communicate the details to your colleagues when appropriate. ## Dependencies * don't introduce any new dependencies without asking. * don't ask for any dependencies that aren't ubiquitous within production environments. ## Commit Messages * use the imperative mood in the subject line. * prefix the subject line with a change type. our change types are: * `documentation`: documentation changes. * `feature`: a new feature. * `refactor`: a code change that neither fixes a bug nor adds a feature. * `bugfix`: a bug fix. * `chore`: routine tasks, maintenance, or tooling changes. * `test`: adding or correcting tests. * restrict the subject line to fifty characters or less. * capitalize the subject line. * do not end the subject line with a period. * separate subject from body with a blank line. ================================================ FILE: TODO.md ================================================ 3. Task cancellation. When API http request gets cancelled, it should cancel corresponding task. 4. I'd like to see profiled network latency / bandwidth. 5. I'd like to see how much bandwidth each link is using. 7. Solve the problem of in continuous batching when a new prompt comes in, it will block decode of the current batch until the prefill is complete. 8. We want people to be able to copy models over to a new device without ever connecting EXO to the internet. Right now EXO require internet connection once to cache some files to check if a download is complete. Instead, we should simply check if there is a non-empty model folder locally with no .partial files. This indicates it's a fully downloaded model that can be loaded. 13. Memory pressure instead of memory used. 14. Show the type of each connection (TB5, Ethernet, etc.) in the UI. Refer to old exo: https://github.com/exo-explore/exo/blob/56f783b38dc6b08ce606b07a5386dc40dae00330/exo/helpers.py#L251 15. Prioritise certain connection types (or by latency). TB5 > Ethernet > WiFi. Refer to old exo: https://github.com/exo-explore/exo/blob/56f783b38dc6b08ce606b07a5386dc40dae00330/exo/helpers.py#L251 16. Dynamically switch to higher priority connection when it becomes available. Probably bring back InstanceReplacedAtomically. 17. Faster model loads by streaming model from other devices in cluster. 18. Add support for specifying the type of network connection to use in a test. Depends on 15/16. 25. Rethink retry logic 27. Log cleanup - per-module log filters and default to DEBUG log levels 28. Validate RDMA connections with ibv_devinfo in the info gatherer ================================================ FILE: app/EXO/EXO/Assets.xcassets/AccentColor.colorset/Contents.json ================================================ { "colors" : [ { "idiom" : "universal" } ], "info" : { "author" : "xcode", "version" : 1 } } ================================================ FILE: app/EXO/EXO/Assets.xcassets/AppIcon.appiconset/Contents.json ================================================ { "images" : [ { "filename" : "16-mac.png", "idiom" : "mac", "scale" : "1x", "size" : "16x16" }, { "filename" : "32-mac.png", "idiom" : "mac", "scale" : "2x", "size" : "16x16" }, { "filename" : "32-mac 1.png", "idiom" : "mac", "scale" : "1x", "size" : "32x32" }, { "filename" : "64-mac.png", "idiom" : "mac", "scale" : "2x", "size" : "32x32" }, { "filename" : "128-mac.png", "idiom" : "mac", "scale" : "1x", "size" : "128x128" }, { "filename" : "256-mac 1.png", "idiom" : "mac", "scale" : "2x", "size" : "128x128" }, { "filename" : "256-mac.png", "idiom" : "mac", "scale" : "1x", "size" : "256x256" }, { "filename" : "512-mac.png", "idiom" : "mac", "scale" : "2x", "size" : "256x256" }, { "filename" : "512-mac 1.png", "idiom" : "mac", "scale" : "1x", "size" : "512x512" }, { "filename" : "1024-mac.png", "idiom" : "mac", "scale" : "2x", "size" : "512x512" } ], "info" : { "author" : "xcode", "version" : 1 } } ================================================ FILE: app/EXO/EXO/Assets.xcassets/Contents.json ================================================ { "info" : { "author" : "xcode", "version" : 1 } } ================================================ FILE: app/EXO/EXO/Assets.xcassets/menubar-icon.imageset/Contents.json ================================================ { "images" : [ { "filename" : "exo-logo-hq-square-transparent-bg.png", "idiom" : "universal", "scale" : "1x" }, { "idiom" : "universal", "scale" : "2x" }, { "idiom" : "universal", "scale" : "3x" } ], "info" : { "author" : "xcode", "version" : 1 } } ================================================ FILE: app/EXO/EXO/ContentView.swift ================================================ // // ContentView.swift // EXO // // Created by Sami Khan on 2025-11-22. // import AppKit import SwiftUI struct ContentView: View { @EnvironmentObject private var controller: ExoProcessController @EnvironmentObject private var stateService: ClusterStateService @EnvironmentObject private var networkStatusService: NetworkStatusService @EnvironmentObject private var localNetworkChecker: LocalNetworkChecker @EnvironmentObject private var updater: SparkleUpdater @EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService @EnvironmentObject private var settingsWindowController: SettingsWindowController @State private var focusedNode: NodeViewModel? @State private var deletingInstanceIDs: Set = [] @State private var showAllNodes = false @State private var showAllInstances = false @State private var baseURLCopied = false @State private var showAdvanced = false @State private var showDebugInfo = false private enum BugReportPhase: Equatable { case idle case prompting case sending(String) case success(String) case failure(String) } @State private var bugReportPhase: BugReportPhase = .idle @State private var bugReportUserDescription: String = "" @State private var uninstallInProgress = false @State private var pendingNamespace: String = "" @State private var pendingHFToken: String = "" @State private var pendingEnableImageModels = false var body: some View { VStack(alignment: .leading, spacing: 12) { statusSection if shouldShowLocalNetworkWarning { localNetworkWarningBanner } if shouldShowClusterDetails { Divider() overviewSection topologySection nodeSection } if shouldShowInstances { instanceSection } Spacer(minLength: 0) controlButtons } .animation(.easeInOut(duration: 0.3), value: shouldShowClusterDetails) .animation(.easeInOut(duration: 0.3), value: shouldShowInstances) .animation(.easeInOut(duration: 0.3), value: shouldShowLocalNetworkWarning) .padding() .frame(width: 340) .onAppear { Task { await networkStatusService.refresh() } } } private var shouldShowLocalNetworkWarning: Bool { // Show warning if local network is not working and EXO is running. // The checker uses a longer timeout on first launch to allow time for // the permission prompt, so this correctly handles both: // 1. User denied permission on first launch // 2. Permission broke after restart (macOS TCC bug) if case .notWorking = localNetworkChecker.status { return controller.status != .stopped } return false } private var localNetworkWarningBanner: some View { VStack(alignment: .leading, spacing: 6) { HStack(spacing: 6) { Image(systemName: "exclamationmark.triangle.fill") .foregroundColor(.orange) Text("Local Network Access Issue") .font(.caption) .fontWeight(.semibold) } Text( "Device discovery won't work. To fix:\n1. Quit EXO\n2. Open System Settings → Privacy & Security → Local Network\n3. Toggle EXO off, then back on\n4. Relaunch EXO" ) .font(.caption2) .foregroundColor(.secondary) .fixedSize(horizontal: false, vertical: true) Button { openLocalNetworkSettings() } label: { Text("Open Settings") .font(.caption2) } .buttonStyle(.bordered) .controlSize(.small) } .padding(8) .background( RoundedRectangle(cornerRadius: 8) .fill(Color.orange.opacity(0.1)) ) .overlay( RoundedRectangle(cornerRadius: 8) .stroke(Color.orange.opacity(0.3), lineWidth: 1) ) } private func openLocalNetworkSettings() { // Open Privacy & Security settings - Local Network section if let url = URL( string: "x-apple.systempreferences:com.apple.preference.security?Privacy_LocalNetwork") { NSWorkspace.shared.open(url) } } private var topologySection: some View { Group { if let topology = stateService.latestSnapshot?.topologyViewModel( localNodeId: stateService.localNodeId), !topology.nodes.isEmpty { TopologyMiniView(topology: topology) } } } private var statusSection: some View { HStack(spacing: 8) { VStack(alignment: .leading, spacing: 2) { Text("EXO") .font(.headline) Text(controller.status.displayText) .font(.caption) .foregroundColor(.secondary) if let detail = statusDetailText { Text(detail) .font(.caption2) .foregroundColor(.secondary) } } Spacer() Toggle("", isOn: processToggleBinding) .toggleStyle(.switch) .labelsHidden() } } private var overviewSection: some View { Group { if let snapshot = stateService.latestSnapshot { let overview = snapshot.overview() VStack(alignment: .leading, spacing: 4) { HStack { VStack(alignment: .leading) { Text( "\(overview.usedRam, specifier: "%.0f") / \(overview.totalRam, specifier: "%.0f") GB" ) .font(.headline) Text("Memory") .font(.caption) .foregroundColor(.secondary) } Spacer() VStack(alignment: .leading) { Text("\(overview.nodeCount)") .font(.headline) Text("Nodes") .font(.caption) .foregroundColor(.secondary) } Spacer() VStack(alignment: .leading) { Text("\(overview.instanceCount)") .font(.headline) Text("Instances") .font(.caption) .foregroundColor(.secondary) } } } } else { Text("Connecting to EXO…") .font(.caption) .foregroundColor(.secondary) } } } private var nodeSection: some View { Group { if let nodes = stateService.latestSnapshot?.nodeViewModels(), !nodes.isEmpty { VStack(alignment: .leading, spacing: 4) { HStack { Text("Nodes") .font(.caption) .foregroundColor(.secondary) Text("(\(nodes.count))") .font(.caption) .foregroundColor(.secondary) Spacer() collapseButton(isExpanded: $showAllNodes) } .animation(nil, value: showAllNodes) if showAllNodes { VStack(alignment: .leading, spacing: 8) { ForEach(nodes) { node in NodeRowView(node: node) .padding(.horizontal, 6) .padding(.vertical, 4) .background(.regularMaterial.opacity(0.6)) .clipShape(RoundedRectangle(cornerRadius: 6)) } } .transition(.opacity) } } .animation(.easeInOut(duration: 0.25), value: showAllNodes) } } } private var instanceSection: some View { Group { if let instances = stateService.latestSnapshot?.instanceViewModels() { VStack(alignment: .leading, spacing: 8) { HStack { Text("Instances") .font(.caption) .foregroundColor(.secondary) Text("(\(instances.count))") .font(.caption) .foregroundColor(.secondary) Spacer() if !instances.isEmpty { collapseButton(isExpanded: $showAllInstances) } } .animation(nil, value: showAllInstances) if showAllInstances, !instances.isEmpty { VStack(alignment: .leading, spacing: 8) { ForEach(instances) { instance in InstanceRowView(instance: instance) .padding(.horizontal, 6) .padding(.vertical, 4) .background(.regularMaterial.opacity(0.6)) .clipShape(RoundedRectangle(cornerRadius: 6)) } } .transition(.opacity) } } .animation(.easeInOut(duration: 0.25), value: showAllInstances) } } } private var controlButtons: some View { VStack(alignment: .leading, spacing: 0) { if controller.status != .stopped { dashboardButton baseURLRow Divider() .padding(.vertical, 8) } else { Divider() .padding(.vertical, 4) } HoverButton( title: "Settings", tint: .primary, trailingSystemImage: "gear" ) { settingsWindowController.open( controller: controller, updater: updater, networkStatusService: networkStatusService, thunderboltBridgeService: thunderboltBridgeService, stateService: stateService ) } HoverButton( title: "Check for Updates", tint: .primary, trailingSystemImage: "arrow.triangle.2.circlepath" ) { updater.checkForUpdates() } .padding(.bottom, 8) HoverButton(title: "Quit", tint: .secondary) { controller.stop() NSApplication.shared.terminate(nil) } } } private var dashboardButton: some View { HoverButton( title: "Web Dashboard", tint: .primary, trailingSystemImage: "arrow.up.right" ) { guard let url = URL(string: "http://localhost:52415/") else { return } NSWorkspace.shared.open(url) } } private var baseURLRow: some View { HStack(spacing: 6) { Image(systemName: "link") .imageScale(.small) .foregroundColor(.secondary) Text("localhost:52415/v1") .font(.system(.caption, design: .monospaced)) .foregroundColor(.primary) Spacer() Button { NSPasteboard.general.clearContents() NSPasteboard.general.setString("http://localhost:52415/v1", forType: .string) baseURLCopied = true DispatchQueue.main.asyncAfter(deadline: .now() + 2) { baseURLCopied = false } } label: { Image(systemName: baseURLCopied ? "checkmark" : "doc.on.doc") .imageScale(.small) .foregroundColor(baseURLCopied ? .green : .secondary) .contentTransition(.symbolEffect(.replace)) } .buttonStyle(.plain) .help("Copy API base URL") } .padding(.vertical, 4) .padding(.horizontal, 8) } private func collapseButton(isExpanded: Binding) -> some View { Button { isExpanded.wrappedValue.toggle() } label: { Label( isExpanded.wrappedValue ? "Hide" : "Show All", systemImage: isExpanded.wrappedValue ? "chevron.up" : "chevron.down" ) .labelStyle(.titleAndIcon) .contentTransition(.symbolEffect(.replace)) } .buttonStyle(.plain) .font(.caption2) } private func instancesToDisplay(_ instances: [InstanceViewModel]) -> [InstanceViewModel] { if showAllInstances { return instances } return [] } private var shouldShowClusterDetails: Bool { controller.status != .stopped } private var shouldShowInstances: Bool { controller.status != .stopped } private var statusDetailText: String? { switch controller.status { case .failed(let message): return message case .stopped: if let countdown = controller.launchCountdownSeconds { return "Launching in \(countdown)s" } return nil default: if let countdown = controller.launchCountdownSeconds { return "Launching in \(countdown)s" } if let lastError = controller.lastError { return lastError } if let message = stateService.lastActionMessage { return message } return nil } } private var thunderboltStatusText: String { switch networkStatusService.status.thunderboltBridgeState { case .some(.disabled): return "Thunderbolt Bridge: Disabled" case .some(.deleted): return "Thunderbolt Bridge: Deleted" case .some(.enabled): return "Thunderbolt Bridge: Enabled" case nil: return "Thunderbolt Bridge: Unknown" } } private var thunderboltStatusColor: Color { switch networkStatusService.status.thunderboltBridgeState { case .some(.disabled), .some(.deleted): return .green case .some(.enabled): return .red case nil: return .secondary } } /// Shows TB bridge status for all nodes from exo cluster state private var clusterThunderboltBridgeView: some View { let bridgeStatuses = stateService.latestSnapshot?.nodeThunderboltBridge ?? [:] let localNodeId = stateService.localNodeId let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:] return VStack(alignment: .leading, spacing: 1) { if bridgeStatuses.isEmpty { Text("Cluster TB Bridge: No data") .font(.caption2) .foregroundColor(.secondary) } else { Text("Cluster TB Bridge Status:") .font(.caption2) .foregroundColor(.secondary) ForEach(Array(bridgeStatuses.keys.sorted()), id: \.self) { nodeId in if let status = bridgeStatuses[nodeId] { let nodeName = nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8)) let isLocal = nodeId == localNodeId let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):" let statusText = !status.exists ? "N/A" : (status.enabled ? "Enabled" : "Disabled") let color: Color = !status.exists ? .secondary : (status.enabled ? .red : .green) Text("\(prefix) \(statusText)") .font(.caption2) .foregroundColor(color) } } } } } private var interfaceIpList: some View { let statuses = networkStatusService.status.interfaceStatuses return VStack(alignment: .leading, spacing: 1) { Text("Interfaces (en0–en7):") .font(.caption2) .foregroundColor(.secondary) if statuses.isEmpty { Text(" Unknown") .font(.caption2) .foregroundColor(.secondary) } else { ForEach(statuses, id: \.interfaceName) { status in let ipText = status.ipAddress ?? "No IP" Text(" \(status.interfaceName): \(ipText)") .font(.caption2) .foregroundColor(status.ipAddress == nil ? .red : .green) } } } } private var debugSection: some View { VStack(alignment: .leading, spacing: 4) { HoverButton( title: "Debug Info", tint: .primary, trailingSystemImage: showDebugInfo ? "chevron.up" : "chevron.down", small: true ) { showDebugInfo.toggle() } if showDebugInfo { VStack(alignment: .leading, spacing: 4) { Text("Version: \(buildTag)") .font(.caption2) .foregroundColor(.secondary) Text("Commit: \(buildCommit)") .font(.caption2) .foregroundColor(.secondary) Text(thunderboltStatusText) .font(.caption2) .foregroundColor(thunderboltStatusColor) clusterThunderboltBridgeView interfaceIpList rdmaStatusView sendBugReportButton .padding(.top, 6) } .padding(.leading, 8) .transition(.opacity) } } .animation(.easeInOut(duration: 0.25), value: showDebugInfo) } private var rdmaStatusView: some View { let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:] let localNodeId = stateService.localNodeId let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:] let localDevices = networkStatusService.status.localRdmaDevices let localPorts = networkStatusService.status.localRdmaActivePorts return VStack(alignment: .leading, spacing: 1) { if rdmaStatuses.isEmpty { Text("Cluster RDMA: No data") .font(.caption2) .foregroundColor(.secondary) } else { Text("Cluster RDMA Status:") .font(.caption2) .foregroundColor(.secondary) ForEach(Array(rdmaStatuses.keys.sorted()), id: \.self) { nodeId in if let status = rdmaStatuses[nodeId] { let nodeName = nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8)) let isLocal = nodeId == localNodeId let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):" let statusText = status.enabled ? "Enabled" : "Disabled" let color: Color = status.enabled ? .green : .orange Text("\(prefix) \(statusText)") .font(.caption2) .foregroundColor(color) } } } if !localDevices.isEmpty { Text(" Local Devices: \(localDevices.joined(separator: ", "))") .font(.caption2) .foregroundColor(.secondary) } if !localPorts.isEmpty { Text(" Local Active Ports:") .font(.caption2) .foregroundColor(.secondary) ForEach(localPorts, id: \.device) { port in Text(" \(port.device) port \(port.port): \(port.state)") .font(.caption2) .foregroundColor(.green) } } } } private var sendBugReportButton: some View { VStack(alignment: .leading, spacing: 6) { switch bugReportPhase { case .idle: Button { bugReportPhase = .prompting bugReportUserDescription = "" } label: { HStack { Text("Send Bug Report") .font(.caption) .fontWeight(.semibold) Spacer() } .padding(.vertical, 6) .padding(.horizontal, 8) .background( RoundedRectangle(cornerRadius: 6) .fill(Color.accentColor.opacity(0.12)) ) } .buttonStyle(.plain) case .prompting: VStack(alignment: .leading, spacing: 6) { Text("What's the issue? (optional)") .font(.caption2) .foregroundColor(.secondary) TextEditor(text: $bugReportUserDescription) .font(.caption2) .frame(height: 60) .overlay( RoundedRectangle(cornerRadius: 4) .stroke(Color.secondary.opacity(0.3), lineWidth: 1) ) HStack(spacing: 8) { Button("Send") { Task { await sendBugReport() } } .font(.caption2) .buttonStyle(.borderedProminent) .controlSize(.small) Button("Cancel") { bugReportPhase = .idle } .font(.caption2) .buttonStyle(.bordered) .controlSize(.small) } } .padding(8) .background( RoundedRectangle(cornerRadius: 6) .fill(Color.accentColor.opacity(0.06)) ) case .sending(let message): HStack(spacing: 6) { ProgressView() .scaleEffect(0.6) Text(message) .font(.caption2) .foregroundColor(.secondary) } case .success(let message): VStack(alignment: .leading, spacing: 6) { Text(message) .font(.caption2) .foregroundColor(.secondary) .fixedSize(horizontal: false, vertical: true) Button { openGitHubIssue() } label: { HStack(spacing: 4) { Image(systemName: "arrow.up.right.square") .imageScale(.small) Text("Create GitHub Issue") .font(.caption2) } } .buttonStyle(.bordered) .controlSize(.small) Button("Done") { bugReportPhase = .idle bugReportUserDescription = "" } .font(.caption2) .buttonStyle(.plain) .foregroundColor(.secondary) } case .failure(let message): VStack(alignment: .leading, spacing: 4) { Text(message) .font(.caption2) .foregroundColor(.red) .fixedSize(horizontal: false, vertical: true) Button("Dismiss") { bugReportPhase = .idle } .font(.caption2) .buttonStyle(.plain) .foregroundColor(.secondary) } } } .animation(.easeInOut(duration: 0.2), value: bugReportPhase) } private var processToggleBinding: Binding { Binding( get: { switch controller.status { case .running, .starting: return true case .stopped, .failed: return false } }, set: { isOn in if isOn { stateService.resetTransientState() stateService.startPolling() controller.cancelPendingLaunch() controller.launchIfNeeded() } else { stateService.stopPolling() controller.stop() stateService.resetTransientState() } } ) } private func bindingForNode(_ node: NodeViewModel) -> Binding { Binding( get: { focusedNode?.id == node.id ? focusedNode : nil }, set: { newValue in if newValue == nil { focusedNode = nil } else { focusedNode = newValue } } ) } private func sendBugReport() async { bugReportPhase = .sending("Collecting logs...") let service = BugReportService() let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines) do { let outcome = try await service.sendReport( isManual: true, userDescription: description.isEmpty ? nil : description ) if outcome.success { bugReportPhase = .success(outcome.message) } else { bugReportPhase = .failure(outcome.message) } } catch { bugReportPhase = .failure(error.localizedDescription) } } private func openGitHubIssue() { let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines) var bodyParts: [String] = [] bodyParts.append("## Describe the bug") bodyParts.append("") if !description.isEmpty { bodyParts.append(description) } else { bodyParts.append("A clear and concise description of what the bug is.") } bodyParts.append("") bodyParts.append("## Environment") bodyParts.append("") bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)") bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))") bodyParts.append("") bodyParts.append("## Additional context") bodyParts.append("") bodyParts.append("A bug report with diagnostic logs was submitted via the app.") let body = bodyParts.joined(separator: "\n") var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")! components.queryItems = [ URLQueryItem(name: "template", value: "bug_report.md"), URLQueryItem(name: "title", value: "[BUG] "), URLQueryItem(name: "body", value: body), URLQueryItem(name: "labels", value: "bug"), ] if let url = components.url { NSWorkspace.shared.open(url) } } private func showUninstallConfirmationAlert() { let alert = NSAlert() alert.messageText = "Uninstall EXO" alert.informativeText = """ This will remove EXO and all its system components: • Network configuration daemon • Launch at login registration • EXO network location The app will be moved to Trash. """ alert.alertStyle = .warning alert.addButton(withTitle: "Uninstall") alert.addButton(withTitle: "Cancel") // Style the Uninstall button as destructive if let uninstallButton = alert.buttons.first { uninstallButton.hasDestructiveAction = true } let response = alert.runModal() if response == .alertFirstButtonReturn { performUninstall() } } private func performUninstall() { uninstallInProgress = true // Stop EXO process first controller.cancelPendingLaunch() controller.stop() stateService.stopPolling() // Run the privileged uninstall on a background thread // Using .utility QoS to avoid priority inversion with NSAppleScript's subprocess DispatchQueue.global(qos: .utility).async { do { // Remove network setup daemon and components (requires admin privileges) try NetworkSetupHelper.uninstall() DispatchQueue.main.async { // Unregister from launch at login LaunchAtLoginHelper.disable() // Move app to trash self.moveAppToTrash() // Quit the app DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) { NSApplication.shared.terminate(nil) } } } catch { DispatchQueue.main.async { self.showErrorAlert(message: error.localizedDescription) self.uninstallInProgress = false } } } } private func showErrorAlert(message: String) { let alert = NSAlert() alert.messageText = "Uninstall Failed" alert.informativeText = message alert.alertStyle = .critical alert.addButton(withTitle: "OK") alert.runModal() } private func moveAppToTrash() { guard let appURL = Bundle.main.bundleURL as URL? else { return } do { try FileManager.default.trashItem(at: appURL, resultingItemURL: nil) } catch { // If we can't trash the app, that's OK - user can do it manually // The important system components have already been cleaned up } } private var buildTag: String { Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown" } private var buildCommit: String { Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown" } } private struct HoverButton: View { let title: String let tint: Color let trailingSystemImage: String? let small: Bool let action: () -> Void init( title: String, tint: Color = .primary, trailingSystemImage: String? = nil, small: Bool = false, action: @escaping () -> Void ) { self.title = title self.tint = tint self.trailingSystemImage = trailingSystemImage self.small = small self.action = action } @State private var isHovering = false var body: some View { Button(action: action) { HStack { Text(title) .font(small ? .caption : nil) Spacer() if let systemName = trailingSystemImage { Image(systemName: systemName) .imageScale(.small) } } .frame(maxWidth: .infinity, alignment: .leading) .padding(.vertical, small ? 4 : 6) .padding(.horizontal, small ? 6 : 8) .background( RoundedRectangle(cornerRadius: 6) .fill( isHovering ? Color.accentColor.opacity(0.1) : Color.clear ) ) } .buttonStyle(.plain) .foregroundColor(tint) .onHover { isHovering = $0 } } } ================================================ FILE: app/EXO/EXO/EXO.entitlements ================================================ com.apple.security.app-sandbox com.apple.security.automation.apple-events com.apple.security.files.user-selected.read-only ================================================ FILE: app/EXO/EXO/EXOApp.swift ================================================ // // EXOApp.swift // EXO // // Created by Sami Khan on 2025-11-22. // import AppKit import CoreImage import CoreImage.CIFilterBuiltins import ServiceManagement import Sparkle import SwiftUI import UserNotifications import os.log struct EXOApp: App { @StateObject private var controller: ExoProcessController @StateObject private var stateService: ClusterStateService @StateObject private var networkStatusService: NetworkStatusService @StateObject private var localNetworkChecker: LocalNetworkChecker @StateObject private var updater: SparkleUpdater @StateObject private var thunderboltBridgeService: ThunderboltBridgeService @StateObject private var settingsWindowController: SettingsWindowController private let terminationObserver: TerminationObserver private let firstLaunchPopout = FirstLaunchPopout() private let ciContext = CIContext(options: nil) init() { let controller = ExoProcessController() let updater = SparkleUpdater(processController: controller) terminationObserver = TerminationObserver { Task { @MainActor in controller.cancelPendingLaunch() controller.stop() } } _controller = StateObject(wrappedValue: controller) let service = ClusterStateService() _stateService = StateObject(wrappedValue: service) let networkStatus = NetworkStatusService() _networkStatusService = StateObject(wrappedValue: networkStatus) let localNetwork = LocalNetworkChecker() _localNetworkChecker = StateObject(wrappedValue: localNetwork) _updater = StateObject(wrappedValue: updater) let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service) _thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge) _settingsWindowController = StateObject(wrappedValue: SettingsWindowController()) enableLaunchAtLoginIfNeeded() // Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops) NetworkSetupHelper.promptAndInstallIfNeeded() // Check local network access periodically (warning disappears when user grants permission) localNetwork.startPeriodicChecking(interval: 10) controller.scheduleLaunch(after: 5) service.startPolling() networkStatus.startPolling() } var body: some Scene { MenuBarExtra { ContentView() .environmentObject(controller) .environmentObject(stateService) .environmentObject(networkStatusService) .environmentObject(localNetworkChecker) .environmentObject(updater) .environmentObject(thunderboltBridgeService) .environmentObject(settingsWindowController) } label: { menuBarIcon .onReceive(controller.$isFirstLaunchReady) { ready in if ready { DispatchQueue.main.asyncAfter(deadline: .now() + 3.0) { self.firstLaunchPopout.onComplete = { [weak controller] in controller?.markOnboardingCompleted() } self.firstLaunchPopout.show() } } } } .menuBarExtraStyle(.window) } private var menuBarIcon: some View { let baseImage = resizedMenuBarIcon(named: "menubar-icon", size: 26) let iconImage: NSImage if controller.status == .stopped, let grey = greyscale(image: baseImage) { iconImage = grey } else { iconImage = baseImage ?? NSImage(named: "menubar-icon") ?? NSImage() } return Image(nsImage: iconImage) .accessibilityLabel("EXO") } private func resizedMenuBarIcon(named: String, size: CGFloat) -> NSImage? { guard let original = NSImage(named: named) else { print("Failed to load image named: \(named)") return nil } let targetSize = NSSize(width: size, height: size) let resized = NSImage(size: targetSize) resized.lockFocus() defer { resized.unlockFocus() } NSGraphicsContext.current?.imageInterpolation = .high original.draw( in: NSRect(origin: .zero, size: targetSize), from: NSRect(origin: .zero, size: original.size), operation: .copy, fraction: 1.0 ) return resized } private func greyscale(image: NSImage?) -> NSImage? { guard let image, let tiff = image.tiffRepresentation, let bitmap = NSBitmapImageRep(data: tiff), let cgImage = bitmap.cgImage else { return nil } let ciImage = CIImage(cgImage: cgImage) let filter = CIFilter.colorControls() filter.inputImage = ciImage filter.saturation = 0 filter.brightness = -0.2 filter.contrast = 0.9 guard let output = filter.outputImage, let rendered = ciContext.createCGImage(output, from: output.extent) else { return nil } return NSImage(cgImage: rendered, size: image.size) } private func enableLaunchAtLoginIfNeeded() { guard SMAppService.mainApp.status != .enabled else { return } do { try SMAppService.mainApp.register() } catch { Logger().error( "Failed to register EXO for launch at login: \(error.localizedDescription)") } } } /// Helper for managing EXO's launch-at-login registration enum LaunchAtLoginHelper { private static let logger = Logger(subsystem: "io.exo.EXO", category: "LaunchAtLogin") /// Unregisters EXO from launching at login static func disable() { guard SMAppService.mainApp.status == .enabled else { return } do { try SMAppService.mainApp.unregister() logger.info("Unregistered EXO from launch at login") } catch { logger.error( "Failed to unregister EXO from launch at login: \(error.localizedDescription, privacy: .public)" ) } } } final class SparkleUpdater: NSObject, ObservableObject { private let controller: SPUStandardUpdaterController private let delegateProxy: ExoUpdaterDelegate private let notificationDelegate = ExoNotificationDelegate() private var periodicCheckTask: Task? init(processController: ExoProcessController) { let proxy = ExoUpdaterDelegate(processController: processController) delegateProxy = proxy controller = SPUStandardUpdaterController( startingUpdater: true, updaterDelegate: proxy, userDriverDelegate: nil ) super.init() let center = UNUserNotificationCenter.current() center.delegate = notificationDelegate center.requestAuthorization(options: [.alert, .sound]) { _, _ in } controller.updater.automaticallyChecksForUpdates = true controller.updater.automaticallyDownloadsUpdates = false controller.updater.updateCheckInterval = 900 // 15 minutes DispatchQueue.main.asyncAfter(deadline: .now() + 5) { [weak controller] in controller?.updater.checkForUpdatesInBackground() } let updater = controller.updater let intervalSeconds = max(60.0, controller.updater.updateCheckInterval) let intervalNanos = UInt64(intervalSeconds * 1_000_000_000) periodicCheckTask = Task { while !Task.isCancelled { try? await Task.sleep(nanoseconds: intervalNanos) await MainActor.run { updater.checkForUpdatesInBackground() } } } } deinit { periodicCheckTask?.cancel() } @MainActor func checkForUpdates() { controller.checkForUpdates(nil) } } private final class ExoUpdaterDelegate: NSObject, SPUUpdaterDelegate { private weak var processController: ExoProcessController? init(processController: ExoProcessController) { self.processController = processController } nonisolated func updater(_ updater: SPUUpdater, didFindValidUpdate item: SUAppcastItem) { showNotification( title: "Update available", body: "EXO \(item.displayVersionString) is ready to install." ) } nonisolated func updaterWillRelaunchApplication(_ updater: SPUUpdater) { Task { @MainActor in guard let controller = self.processController else { return } controller.cancelPendingLaunch() controller.stop() } } nonisolated private func showNotification(title: String, body: String) { let center = UNUserNotificationCenter.current() let content = UNMutableNotificationContent() content.title = title content.body = body let request = UNNotificationRequest( identifier: "exo-update-\(UUID().uuidString)", content: content, trigger: nil ) center.add(request, withCompletionHandler: nil) } } private final class ExoNotificationDelegate: NSObject, UNUserNotificationCenterDelegate { func userNotificationCenter( _ center: UNUserNotificationCenter, willPresent notification: UNNotification, withCompletionHandler completionHandler: @escaping (UNNotificationPresentationOptions) -> Void ) { completionHandler([.banner, .list, .sound]) } } @MainActor private final class TerminationObserver { private var token: NSObjectProtocol? init(onTerminate: @escaping () -> Void) { token = NotificationCenter.default.addObserver( forName: NSApplication.willTerminateNotification, object: nil, queue: .main ) { _ in onTerminate() } } deinit { if let token { NotificationCenter.default.removeObserver(token) } } } ================================================ FILE: app/EXO/EXO/ExoProcessController.swift ================================================ import AppKit import Combine import Foundation private let customNamespaceKey = "EXOCustomNamespace" private let hfTokenKey = "EXOHFToken" private let enableImageModelsKey = "EXOEnableImageModels" private let offlineModeKey = "EXOOfflineMode" private let onboardingCompletedKey = "EXOOnboardingCompleted" @MainActor final class ExoProcessController: ObservableObject { enum Status: Equatable { case stopped case starting case running case failed(message: String) var displayText: String { switch self { case .stopped: return "Stopped" case .starting: return "Starting…" case .running: return "Running" case .failed: return "Failed" } } } static let exoDirectoryURL: URL = { URL(fileURLWithPath: NSHomeDirectory()).appendingPathComponent(".exo") }() @Published private(set) var status: Status = .stopped @Published private(set) var lastError: String? @Published private(set) var launchCountdownSeconds: Int? @Published var customNamespace: String = { return UserDefaults.standard.string(forKey: customNamespaceKey) ?? "" }() { didSet { UserDefaults.standard.set(customNamespace, forKey: customNamespaceKey) } } @Published var hfToken: String = { return UserDefaults.standard.string(forKey: hfTokenKey) ?? "" }() { didSet { UserDefaults.standard.set(hfToken, forKey: hfTokenKey) } } @Published var enableImageModels: Bool = { return UserDefaults.standard.bool(forKey: enableImageModelsKey) }() { didSet { UserDefaults.standard.set(enableImageModels, forKey: enableImageModelsKey) } } @Published var offlineMode: Bool = { return UserDefaults.standard.bool(forKey: offlineModeKey) }() { didSet { UserDefaults.standard.set(offlineMode, forKey: offlineModeKey) } } /// Fires once when EXO transitions to `.running` for the very first time (fresh install). @Published private(set) var isFirstLaunchReady = false private var process: Process? private var runtimeDirectoryURL: URL? private var pendingLaunchTask: Task? func launchIfNeeded() { guard process?.isRunning != true else { return } launch() } func launch() { do { guard process?.isRunning != true else { return } cancelPendingLaunch() status = .starting lastError = nil let runtimeURL = try resolveRuntimeDirectory() runtimeDirectoryURL = runtimeURL let executableURL = runtimeURL.appendingPathComponent("exo") let child = Process() child.executableURL = executableURL let exoHomeURL = Self.exoDirectoryURL try? FileManager.default.createDirectory( at: exoHomeURL, withIntermediateDirectories: true ) child.currentDirectoryURL = exoHomeURL child.environment = makeEnvironment(for: runtimeURL) child.standardOutput = FileHandle.nullDevice child.standardError = FileHandle.nullDevice child.terminationHandler = { [weak self] proc in Task { @MainActor in guard let self else { return } self.process = nil switch self.status { case .stopped: break case .failed: break default: self.status = .failed( message: "Exited with code \(proc.terminationStatus)" ) self.lastError = "Process exited with code \(proc.terminationStatus)" } } } try child.run() process = child status = .running // Show welcome popout on every launch isFirstLaunchReady = true } catch { process = nil status = .failed(message: "Launch error") lastError = error.localizedDescription } } func stop() { guard let process else { status = .stopped return } process.terminationHandler = nil status = .stopped guard process.isRunning else { self.process = nil return } let proc = process self.process = nil Task.detached { proc.interrupt() for _ in 0..<50 { if !proc.isRunning { return } try? await Task.sleep(nanoseconds: 100_000_000) } if proc.isRunning { proc.terminate() } for _ in 0..<30 { if !proc.isRunning { return } try? await Task.sleep(nanoseconds: 100_000_000) } if proc.isRunning { kill(proc.processIdentifier, SIGKILL) } } } func restart() { stop() launch() } /// Mark onboarding as completed (user interacted with the welcome popout). func markOnboardingCompleted() { UserDefaults.standard.set(true, forKey: onboardingCompletedKey) } /// Reset onboarding so the welcome popout appears on next launch. func resetOnboarding() { UserDefaults.standard.removeObject(forKey: onboardingCompletedKey) isFirstLaunchReady = false } func scheduleLaunch(after seconds: TimeInterval) { cancelPendingLaunch() let start = max(1, Int(ceil(seconds))) pendingLaunchTask = Task { [weak self] in guard let self else { return } await MainActor.run { self.launchCountdownSeconds = start } var remaining = start while remaining > 0 { try? await Task.sleep(nanoseconds: 1_000_000_000) remaining -= 1 if Task.isCancelled { return } await MainActor.run { if remaining > 0 { self.launchCountdownSeconds = remaining } else { self.launchCountdownSeconds = nil self.launchIfNeeded() } } } } } func cancelPendingLaunch() { pendingLaunchTask?.cancel() pendingLaunchTask = nil launchCountdownSeconds = nil } func revealRuntimeDirectory() { guard let runtimeDirectoryURL else { return } NSWorkspace.shared.activateFileViewerSelecting([runtimeDirectoryURL]) } func statusTintColor() -> NSColor { switch status { case .running: return .systemGreen case .starting: return .systemYellow case .failed: return .systemRed case .stopped: return .systemGray } } private func resolveRuntimeDirectory() throws -> URL { let fileManager = FileManager.default if let override = ProcessInfo.processInfo.environment["EXO_RUNTIME_DIR"] { let url = URL(fileURLWithPath: override).standardizedFileURL if fileManager.fileExists(atPath: url.path) { return url } } if let resourceRoot = Bundle.main.resourceURL { let bundled = resourceRoot.appendingPathComponent("exo", isDirectory: true) if fileManager.fileExists(atPath: bundled.path) { return bundled } } let repoCandidate = URL(fileURLWithPath: fileManager.currentDirectoryPath) .appendingPathComponent("dist/exo", isDirectory: true) if fileManager.fileExists(atPath: repoCandidate.path) { return repoCandidate } throw RuntimeError("Unable to locate the packaged EXO runtime.") } private func makeEnvironment(for runtimeURL: URL) -> [String: String] { var environment = ProcessInfo.processInfo.environment environment["EXO_RUNTIME_DIR"] = runtimeURL.path environment["EXO_LIBP2P_NAMESPACE"] = computeNamespace() if !hfToken.isEmpty { environment["HF_TOKEN"] = hfToken } if enableImageModels { environment["EXO_ENABLE_IMAGE_MODELS"] = "true" } if offlineMode { environment["EXO_OFFLINE"] = "true" } var paths: [String] = [] if let existing = environment["PATH"], !existing.isEmpty { paths = existing.split(separator: ":").map(String.init) } let required = [ runtimeURL.path, runtimeURL.appendingPathComponent("_internal").path, "/opt/homebrew/bin", "/usr/local/bin", "/usr/bin", "/bin", "/usr/sbin", "/sbin", ] for entry in required.reversed() { if !paths.contains(entry) { paths.insert(entry, at: 0) } } environment["PATH"] = paths.joined(separator: ":") return environment } private func buildTag() -> String { if let tag = Bundle.main.infoDictionary?["EXOBuildTag"] as? String, !tag.isEmpty { return tag } if let short = Bundle.main.infoDictionary?["CFBundleShortVersionString"] as? String, !short.isEmpty { return short } return "dev" } private func computeNamespace() -> String { let base = buildTag() let custom = customNamespace.trimmingCharacters(in: .whitespaces) return custom.isEmpty ? base : custom } } struct RuntimeError: LocalizedError { let message: String init(_ message: String) { self.message = message } var errorDescription: String? { message } } ================================================ FILE: app/EXO/EXO/Info.plist ================================================ SUFeedURL https://assets.exolabs.net/appcast.xml EXOBuildTag $(EXO_BUILD_TAG) EXOBuildCommit $(EXO_BUILD_COMMIT) EXOBugReportPresignedUrlEndpoint $(EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT) NSLocalNetworkUsageDescription EXO needs local network access to discover and connect to other devices in your cluster for distributed AI inference. NSBonjourServices _p2p._tcp _p2p._udp _libp2p._udp ================================================ FILE: app/EXO/EXO/Models/ClusterState.swift ================================================ import Foundation // MARK: - API payloads struct ClusterState: Decodable { let instances: [String: ClusterInstance] let runners: [String: RunnerStatusSummary] let tasks: [String: ClusterTask] let topology: Topology? let downloads: [String: [NodeDownloadStatus]] let thunderboltBridgeCycles: [[String]] // Granular node state (split from the old nodeProfiles) let nodeIdentities: [String: NodeIdentity] let nodeMemory: [String: MemoryInfo] let nodeSystem: [String: SystemInfo] let nodeThunderboltBridge: [String: ThunderboltBridgeStatus] let nodeRdmaCtl: [String: NodeRdmaCtlStatus] /// Computed property for backwards compatibility - merges granular state into NodeProfile var nodeProfiles: [String: NodeProfile] { var profiles: [String: NodeProfile] = [:] let allNodeIds = Set(nodeIdentities.keys) .union(nodeMemory.keys) .union(nodeSystem.keys) for nodeId in allNodeIds { let identity = nodeIdentities[nodeId] let memory = nodeMemory[nodeId] let system = nodeSystem[nodeId] profiles[nodeId] = NodeProfile( modelId: identity?.modelId, chipId: identity?.chipId, friendlyName: identity?.friendlyName, memory: memory, system: system ) } return profiles } init(from decoder: Decoder) throws { let container = try decoder.container(keyedBy: CodingKeys.self) let rawInstances = try container.decode([String: TaggedInstance].self, forKey: .instances) self.instances = rawInstances.mapValues(\.instance) self.runners = try container.decode([String: RunnerStatusSummary].self, forKey: .runners) let rawTasks = try container.decodeIfPresent([String: TaggedTask].self, forKey: .tasks) ?? [:] self.tasks = rawTasks.compactMapValues(\.task) self.topology = try container.decodeIfPresent(Topology.self, forKey: .topology) let rawDownloads = try container.decodeIfPresent([String: [TaggedNodeDownload]].self, forKey: .downloads) ?? [:] self.downloads = rawDownloads.mapValues { $0.compactMap(\.status) } self.thunderboltBridgeCycles = try container.decodeIfPresent([[String]].self, forKey: .thunderboltBridgeCycles) ?? [] // Granular node state self.nodeIdentities = try container.decodeIfPresent([String: NodeIdentity].self, forKey: .nodeIdentities) ?? [:] self.nodeMemory = try container.decodeIfPresent([String: MemoryInfo].self, forKey: .nodeMemory) ?? [:] self.nodeSystem = try container.decodeIfPresent([String: SystemInfo].self, forKey: .nodeSystem) ?? [:] self.nodeThunderboltBridge = try container.decodeIfPresent( [String: ThunderboltBridgeStatus].self, forKey: .nodeThunderboltBridge ) ?? [:] self.nodeRdmaCtl = try container.decodeIfPresent( [String: NodeRdmaCtlStatus].self, forKey: .nodeRdmaCtl ) ?? [:] } private enum CodingKeys: String, CodingKey { case instances case runners case topology case tasks case downloads case thunderboltBridgeCycles case nodeIdentities case nodeMemory case nodeSystem case nodeThunderboltBridge case nodeRdmaCtl } } private struct TaggedInstance: Decodable { let instance: ClusterInstance init(from decoder: Decoder) throws { let container = try decoder.singleValueContainer() let payloads = try container.decode([String: ClusterInstancePayload].self) guard let entry = payloads.first else { throw DecodingError.dataCorrupted( DecodingError.Context( codingPath: decoder.codingPath, debugDescription: "Empty instance payload") ) } self.instance = ClusterInstance( instanceId: entry.value.instanceId, shardAssignments: entry.value.shardAssignments, sharding: entry.key ) } } private struct ClusterInstancePayload: Decodable { let instanceId: String? let shardAssignments: ShardAssignments } struct ClusterInstance { let instanceId: String? let shardAssignments: ShardAssignments let sharding: String } struct ShardAssignments: Decodable { let modelId: String let nodeToRunner: [String: String] } struct RunnerStatusSummary: Decodable { let status: String let errorMessage: String? init(from decoder: Decoder) throws { let container = try decoder.singleValueContainer() let payloads = try container.decode([String: RunnerStatusDetail].self) guard let entry = payloads.first else { throw DecodingError.dataCorrupted( DecodingError.Context( codingPath: decoder.codingPath, debugDescription: "Empty runner status payload") ) } self.status = entry.key self.errorMessage = entry.value.errorMessage } } struct RunnerStatusDetail: Decodable { let errorMessage: String? } struct NodeProfile: Decodable { let modelId: String? let chipId: String? let friendlyName: String? let memory: MemoryInfo? let system: SystemInfo? } struct NodeIdentity: Decodable { let modelId: String? let chipId: String? let friendlyName: String? } struct ThunderboltBridgeStatus: Decodable { let enabled: Bool let exists: Bool let serviceName: String? } struct NodeRdmaCtlStatus: Decodable { let enabled: Bool } struct MemoryInfo: Decodable { let ramTotal: MemoryValue? let ramAvailable: MemoryValue? } struct MemoryValue: Decodable { let inBytes: Int64? } struct SystemInfo: Decodable { let gpuUsage: Double? let temp: Double? let sysPower: Double? let pcpuUsage: Double? let ecpuUsage: Double? } struct Topology: Decodable { /// Node IDs in the topology let nodes: [String] /// Flattened list of connections (source -> sink pairs) let connections: [TopologyConnection] init(from decoder: Decoder) throws { let container = try decoder.container(keyedBy: CodingKeys.self) self.nodes = try container.decodeIfPresent([String].self, forKey: .nodes) ?? [] // Connections come as nested map: { source: { sink: [edges] } } // We flatten to array of (source, sink) pairs var flatConnections: [TopologyConnection] = [] if let nested = try container.decodeIfPresent( [String: [String: [AnyCodable]]].self, forKey: .connections ) { for (source, sinks) in nested { for sink in sinks.keys { flatConnections.append( TopologyConnection(localNodeId: source, sendBackNodeId: sink)) } } } self.connections = flatConnections } private enum CodingKeys: String, CodingKey { case nodes case connections } } /// Placeholder for decoding arbitrary JSON values we don't need to inspect private struct AnyCodable: Decodable { init(from decoder: Decoder) throws { // Just consume the value without storing it _ = try? decoder.singleValueContainer().decode(Bool.self) _ = try? decoder.singleValueContainer().decode(Int.self) _ = try? decoder.singleValueContainer().decode(Double.self) _ = try? decoder.singleValueContainer().decode(String.self) _ = try? decoder.singleValueContainer().decode([AnyCodable].self) _ = try? decoder.singleValueContainer().decode([String: AnyCodable].self) } } struct TopologyConnection { let localNodeId: String let sendBackNodeId: String } // MARK: - Downloads private struct TaggedNodeDownload: Decodable { let status: NodeDownloadStatus? init(from decoder: Decoder) throws { let container = try decoder.singleValueContainer() let payloads = try container.decode([String: NodeDownloadPayload].self) guard let entry = payloads.first else { status = nil return } status = NodeDownloadStatus(statusKey: entry.key, payload: entry.value) } } struct NodeDownloadPayload: Decodable { let nodeId: String? let downloadProgress: DownloadProgress? } struct NodeDownloadStatus { let nodeId: String let progress: DownloadProgress? init?(statusKey: String, payload: NodeDownloadPayload) { guard let nodeId = payload.nodeId else { return nil } self.nodeId = nodeId self.progress = statusKey == "DownloadOngoing" ? payload.downloadProgress : nil } } struct DownloadProgress: Decodable { let totalBytes: ByteValue let downloadedBytes: ByteValue let speed: Double? let etaMs: Int64? let completedFiles: Int? let totalFiles: Int? let files: [String: FileDownloadProgress]? } struct ByteValue: Decodable { let inBytes: Int64 } struct FileDownloadProgress: Decodable { let totalBytes: ByteValue let downloadedBytes: ByteValue let speed: Double? let etaMs: Int64? } // MARK: - Tasks struct ClusterTask { enum Kind { case chatCompletion } let id: String let status: TaskStatus let instanceId: String? let kind: Kind let modelName: String? let promptPreview: String? let errorMessage: String? let parameters: TextGenerationTaskParameters? var sortPriority: Int { switch status { case .running: return 0 case .pending: return 1 case .complete: return 2 case .failed: return 3 case .unknown: return 4 } } } private struct TaggedTask: Decodable { let task: ClusterTask? init(from decoder: Decoder) throws { let container = try decoder.singleValueContainer() let payloads = try container.decode([String: ClusterTaskPayload].self) guard let entry = payloads.first else { task = nil return } task = ClusterTask(kindKey: entry.key, payload: entry.value) } } struct ClusterTaskPayload: Decodable { let taskId: String? let taskStatus: TaskStatus? let instanceId: String? let commandId: String? let taskParams: TextGenerationTaskParameters? let errorType: String? let errorMessage: String? } struct TextGenerationTaskParameters: Decodable, Equatable { let model: String? let messages: [ChatCompletionMessage]? let maxTokens: Int? let stream: Bool? let temperature: Double? let topP: Double? private enum CodingKeys: String, CodingKey { case model case messages case maxTokens case stream case temperature case topP } func promptPreview() -> String? { guard let messages else { return nil } if let userMessage = messages.last(where: { $0.role?.lowercased() == "user" && ($0.content?.isEmpty == false) }) { return userMessage.content } return messages.last?.content } } struct ChatCompletionMessage: Decodable, Equatable { let role: String? let content: String? } extension ClusterTask { init?(kindKey: String, payload: ClusterTaskPayload) { guard let id = payload.taskId else { return nil } let status = payload.taskStatus ?? .unknown switch kindKey { case "TextGeneration": self.init( id: id, status: status, instanceId: payload.instanceId, kind: .chatCompletion, modelName: payload.taskParams?.model, promptPreview: payload.taskParams?.promptPreview(), errorMessage: payload.errorMessage, parameters: payload.taskParams ) default: return nil } } } enum TaskStatus: String, Decodable { case pending = "Pending" case running = "Running" case complete = "Complete" case failed = "Failed" case unknown init(from decoder: Decoder) throws { let container = try decoder.singleValueContainer() let value = try container.decode(String.self) self = TaskStatus(rawValue: value) ?? .unknown } var displayLabel: String { switch self { case .pending, .running, .complete, .failed: return rawValue case .unknown: return "Unknown" } } } // MARK: - Derived summaries struct ClusterOverview { let totalRam: Double let usedRam: Double let nodeCount: Int let instanceCount: Int } struct NodeSummary: Identifiable { let id: String let friendlyName: String let model: String let usedRamGB: Double let totalRamGB: Double let gpuUsagePercent: Double let temperatureCelsius: Double } struct InstanceSummary: Identifiable { let id: String let modelId: String let nodeCount: Int let statusText: String } extension ClusterState { func overview() -> ClusterOverview { var total: Double = 0 var available: Double = 0 for profile in nodeProfiles.values { if let totalBytes = profile.memory?.ramTotal?.inBytes { total += Double(totalBytes) } if let availableBytes = profile.memory?.ramAvailable?.inBytes { available += Double(availableBytes) } } let totalGB = total / 1_073_741_824.0 let usedGB = max(total - available, 0) / 1_073_741_824.0 return ClusterOverview( totalRam: totalGB, usedRam: usedGB, nodeCount: nodeProfiles.count, instanceCount: instances.count ) } func availableModels() -> [ModelOption] { [] } } ================================================ FILE: app/EXO/EXO/Preview Content/Preview Assets.xcassets/Contents.json ================================================ { "info" : { "author" : "xcode", "version" : 1 } } ================================================ FILE: app/EXO/EXO/Services/BugReportService.swift ================================================ import Foundation struct BugReportOutcome: Equatable { let success: Bool let message: String } enum BugReportError: LocalizedError { case invalidEndpoint case presignedUrlFailed(String) case uploadFailed(String) case collectFailed(String) var errorDescription: String? { switch self { case .invalidEndpoint: return "Bug report endpoint is invalid." case .presignedUrlFailed(let message): return "Failed to get presigned URLs: \(message)" case .uploadFailed(let message): return "Bug report upload failed: \(message)" case .collectFailed(let message): return "Bug report collection failed: \(message)" } } } struct BugReportService { private struct PresignedUrlsRequest: Codable { let keys: [String] } private struct PresignedUrlsResponse: Codable { let urls: [String: String] let expiresIn: Int? } func sendReport( baseURL: URL = URL(string: "http://127.0.0.1:52415")!, now: Date = Date(), isManual: Bool = false, userDescription: String? = nil ) async throws -> BugReportOutcome { let timestamp = Self.runTimestampString(now) let dayPrefix = Self.dayPrefixString(now) let prefix = "reports/\(dayPrefix)/\(timestamp)/" let logFiles = readAllLogs() let ifconfigText = try await captureIfconfig() let hostName = Host.current().localizedName ?? "unknown" let debugInfo = readDebugInfo() let stateData = try await fetch(url: baseURL.appendingPathComponent("state")) // Extract cluster TB bridge status from exo state let clusterTbBridgeStatus = extractClusterTbBridgeStatus(from: stateData) let reportJSON = makeReportJson( timestamp: timestamp, hostName: hostName, ifconfig: ifconfigText, debugInfo: debugInfo, isManual: isManual, clusterTbBridgeStatus: clusterTbBridgeStatus, userDescription: userDescription ) let eventLogFiles = readAllEventLogs() var uploads: [(path: String, data: Data?)] = logFiles.map { (path, data) in ("\(prefix)\(path)", data) } uploads.append( contentsOf: eventLogFiles.map { (path, data) in ("\(prefix)\(path)", data as Data?) }) uploads.append(contentsOf: [ ("\(prefix)state.json", stateData), ("\(prefix)report.json", reportJSON), ]) let uploadItems: [(key: String, body: Data)] = uploads.compactMap { item in guard let body = item.data else { return nil } return (key: item.path, body: body) } guard !uploadItems.isEmpty else { return BugReportOutcome(success: false, message: "No data to upload") } let presignedUrls = try await fetchPresignedUploadUrls(keys: uploadItems.map(\.key)) for item in uploadItems { guard let urlString = presignedUrls[item.key], let url = URL(string: urlString) else { throw BugReportError.uploadFailed("Missing presigned URL for \(item.key)") } try await uploadToPresignedUrl(url: url, body: item.body) } return BugReportOutcome( success: true, message: "Bug Report sent. Thank you for helping to improve EXO 1.0.") } private static func dayPrefixString(_ date: Date) -> String { var calendar = Calendar(identifier: .gregorian) calendar.timeZone = TimeZone(secondsFromGMT: 0) ?? .current let components = calendar.dateComponents([.year, .month, .day], from: date) let year = components.year ?? 0 let month = components.month ?? 0 let day = components.day ?? 0 return String(format: "%04d/%02d/%02d", year, month, day) } private static func runTimestampString(_ date: Date) -> String { let formatter = DateFormatter() formatter.locale = Locale(identifier: "en_US_POSIX") formatter.timeZone = TimeZone(secondsFromGMT: 0) ?? .current formatter.dateFormat = "yyyy-MM-dd'T'HHmmss.SSS'Z'" return formatter.string(from: date) } private func fetchPresignedUploadUrls(keys: [String], bundle: Bundle = .main) async throws -> [String: String] { guard let endpointString = bundle.infoDictionary?["EXOBugReportPresignedUrlEndpoint"] as? String else { throw BugReportError.invalidEndpoint } let trimmedEndpointString = endpointString.trimmingCharacters(in: .whitespacesAndNewlines) guard !trimmedEndpointString.isEmpty, let endpoint = URL(string: trimmedEndpointString) else { throw BugReportError.invalidEndpoint } var request = URLRequest(url: endpoint) request.httpMethod = "POST" request.timeoutInterval = 10 request.setValue("application/json", forHTTPHeaderField: "Content-Type") let encoder = JSONEncoder() request.httpBody = try encoder.encode(PresignedUrlsRequest(keys: keys)) let (data, response) = try await URLSession.shared.data(for: request) guard let http = response as? HTTPURLResponse else { throw BugReportError.presignedUrlFailed("Non-HTTP response") } guard (200..<300).contains(http.statusCode) else { throw BugReportError.presignedUrlFailed("HTTP status \(http.statusCode)") } let decoder = JSONDecoder() let decoded = try decoder.decode(PresignedUrlsResponse.self, from: data) return decoded.urls } private func readAllLogs() -> [(path: String, data: Data)] { let dir = URL(fileURLWithPath: NSHomeDirectory()) .appendingPathComponent(".exo") .appendingPathComponent("exo_log") var results: [(path: String, data: Data)] = [] let contents = (try? FileManager.default.contentsOfDirectory(atPath: dir.path)) ?? [] for name in contents { if let data = try? Data(contentsOf: dir.appendingPathComponent(name)) { results.append(("exo_log/\(name)", data)) } } return results } private func readAllEventLogs() -> [(path: String, data: Data)] { let eventLogDir = URL(fileURLWithPath: NSHomeDirectory()) .appendingPathComponent(".exo") .appendingPathComponent("event_log") var results: [(path: String, data: Data)] = [] for subdir in ["master", "api"] { let dir = eventLogDir.appendingPathComponent(subdir) let contents = (try? FileManager.default.contentsOfDirectory(atPath: dir.path)) ?? [] for name in contents where name.hasPrefix("events.") { if let data = try? Data(contentsOf: dir.appendingPathComponent(name)) { results.append(("event_log/\(subdir)/\(name)", data)) } } } return results } private func captureIfconfig() async throws -> String { let result = runCommand(["/sbin/ifconfig"]) guard result.exitCode == 0 else { throw BugReportError.collectFailed( result.error.isEmpty ? "ifconfig failed" : result.error) } return result.output } private func readDebugInfo() -> DebugInfo { DebugInfo( thunderboltBridgeDisabled: readThunderboltBridgeDisabled(), interfaces: readInterfaces(), rdma: readRDMADebugInfo() ) } private func readRDMADebugInfo() -> DebugInfo.RDMADebugInfo { DebugInfo.RDMADebugInfo( rdmaCtlStatus: safeRunCommand(["/usr/bin/rdma_ctl", "status"]), ibvDevices: safeRunCommand(["/usr/bin/ibv_devices"]), ibvDevinfo: safeRunCommand(["/usr/bin/ibv_devinfo"]) ) } private func readThunderboltBridgeDisabled() -> Bool? { // Dynamically find the Thunderbolt Bridge service (don't assume the name) guard let serviceName = ThunderboltBridgeDetector.findThunderboltBridgeServiceName() else { // No bridge containing Thunderbolt interfaces exists return nil } guard let isEnabled = ThunderboltBridgeDetector.isServiceEnabled(serviceName: serviceName) else { return nil } // Return true if disabled, false if enabled return !isEnabled } private func readInterfaces() -> [DebugInfo.InterfaceStatus] { (0...7).map { "en\($0)" }.map { iface in let result = runCommand(["/sbin/ifconfig", iface]) guard result.exitCode == 0 else { return DebugInfo.InterfaceStatus(name: iface, ip: nil) } let ip = firstInet(from: result.output) return DebugInfo.InterfaceStatus(name: iface, ip: ip) } } private func firstInet(from ifconfigOutput: String) -> String? { for line in ifconfigOutput.split(separator: "\n") { let trimmed = line.trimmingCharacters(in: .whitespaces) guard trimmed.hasPrefix("inet ") else { continue } let parts = trimmed.split(separator: " ") if parts.count >= 2 { let candidate = String(parts[1]) if candidate != "127.0.0.1" { return candidate } } } return nil } private func fetch(url: URL) async throws -> Data? { var request = URLRequest(url: url) request.timeoutInterval = 5 do { let (data, response) = try await URLSession.shared.data(for: request) guard let http = response as? HTTPURLResponse, (200..<300).contains(http.statusCode) else { return nil } return data } catch { return nil } } private func uploadToPresignedUrl(url: URL, body: Data) async throws { let maxAttempts = 2 var lastError: Error? for attempt in 1...maxAttempts { do { var request = URLRequest(url: url) request.httpMethod = "PUT" request.httpBody = body request.timeoutInterval = 30 let (_, response) = try await URLSession.shared.data(for: request) guard let http = response as? HTTPURLResponse else { throw BugReportError.uploadFailed("Non-HTTP response") } guard (200..<300).contains(http.statusCode) else { throw BugReportError.uploadFailed("HTTP status \(http.statusCode)") } return } catch { lastError = error if attempt < maxAttempts { try await Task.sleep(nanoseconds: 400_000_000) } } } throw BugReportError.uploadFailed(lastError?.localizedDescription ?? "Unknown error") } private func makeReportJson( timestamp: String, hostName: String, ifconfig: String, debugInfo: DebugInfo, isManual: Bool, clusterTbBridgeStatus: [[String: Any]]?, userDescription: String? = nil ) -> Data? { let system = readSystemMetadata() let exo = readExoMetadata() var payload: [String: Any] = [ "timestamp": timestamp, "host": hostName, "ifconfig": ifconfig, "debug": debugInfo.toDictionary(), "system": system, "exo_version": exo.version as Any, "exo_commit": exo.commit as Any, "report_type": isManual ? "manual" : "automated", ] if let tbStatus = clusterTbBridgeStatus { payload["cluster_thunderbolt_bridge"] = tbStatus } if let desc = userDescription, !desc.isEmpty { payload["user_description"] = desc } return try? JSONSerialization.data(withJSONObject: payload, options: [.prettyPrinted]) } /// Extracts cluster-wide Thunderbolt Bridge status from exo state JSON private func extractClusterTbBridgeStatus(from stateData: Data?) -> [[String: Any]]? { guard let data = stateData, let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], let nodeThunderboltBridge = json["node_thunderbolt_bridge"] as? [String: [String: Any]] else { return nil } var result: [[String: Any]] = [] for (nodeId, status) in nodeThunderboltBridge { var entry: [String: Any] = ["node_id": nodeId] if let enabled = status["enabled"] as? Bool { entry["enabled"] = enabled } if let exists = status["exists"] as? Bool { entry["exists"] = exists } if let serviceName = status["service_name"] as? String { entry["service_name"] = serviceName } result.append(entry) } return result.isEmpty ? nil : result } private func readSystemMetadata() -> [String: Any] { let hostname = safeRunCommand(["/bin/hostname"]) let computerName = safeRunCommand(["/usr/sbin/scutil", "--get", "ComputerName"]) let localHostName = safeRunCommand(["/usr/sbin/scutil", "--get", "LocalHostName"]) let hostNameCommand = safeRunCommand(["/usr/sbin/scutil", "--get", "HostName"]) let hardwareModel = safeRunCommand(["/usr/sbin/sysctl", "-n", "hw.model"]) let hardwareProfile = safeRunCommand(["/usr/sbin/system_profiler", "SPHardwareDataType"]) let hardwareUUID = hardwareProfile.flatMap(extractHardwareUUID) let osVersion = safeRunCommand(["/usr/bin/sw_vers", "-productVersion"]) let osBuild = safeRunCommand(["/usr/bin/sw_vers", "-buildVersion"]) let kernel = safeRunCommand(["/usr/bin/uname", "-srv"]) let arch = safeRunCommand(["/usr/bin/uname", "-m"]) let routeInfo = safeRunCommand(["/sbin/route", "-n", "get", "default"]) let defaultInterface = routeInfo.flatMap(parseDefaultInterface) let defaultIP = defaultInterface.flatMap { iface in safeRunCommand(["/usr/sbin/ipconfig", "getifaddr", iface]) } let defaultMac = defaultInterface.flatMap { iface in safeRunCommand(["/sbin/ifconfig", iface]).flatMap(parseEtherAddress) } let user = safeRunCommand(["/usr/bin/whoami"]) let consoleUser = safeRunCommand(["/usr/bin/stat", "-f%Su", "/dev/console"]) let uptime = safeRunCommand(["/usr/bin/uptime"]) let diskRoot = safeRunCommand([ "/bin/sh", "-c", "/bin/df -h / | awk 'NR==2 {print $1, $2, $3, $4, $5}'", ]) let interfacesList = safeRunCommand(["/usr/sbin/ipconfig", "getiflist"]) let interfacesAndIPs = interfacesList? .split(whereSeparator: { $0 == " " || $0 == "\n" }) .compactMap { iface -> [String: Any]? in let name = String(iface) guard let ip = safeRunCommand(["/usr/sbin/ipconfig", "getifaddr", name]) else { return nil } return ["name": name, "ip": ip] } ?? [] let wifiSSID: String? let airportPath = "/System/Library/PrivateFrameworks/Apple80211.framework/Versions/Current/Resources/airport" if FileManager.default.isExecutableFile(atPath: airportPath) { wifiSSID = safeRunCommand([airportPath, "-I"]).flatMap(parseWifiSSID) } else { wifiSSID = nil } return [ "hostname": hostname as Any, "computer_name": computerName as Any, "local_hostname": localHostName as Any, "host_name": hostNameCommand as Any, "hardware_model": hardwareModel as Any, "hardware_profile": hardwareProfile as Any, "hardware_uuid": hardwareUUID as Any, "os_version": osVersion as Any, "os_build": osBuild as Any, "kernel": kernel as Any, "arch": arch as Any, "default_interface": defaultInterface as Any, "default_ip": defaultIP as Any, "default_mac": defaultMac as Any, "user": user as Any, "console_user": consoleUser as Any, "uptime": uptime as Any, "disk_root": diskRoot as Any, "interfaces_and_ips": interfacesAndIPs, "ipconfig_getiflist": interfacesList as Any, "wifi_ssid": wifiSSID as Any, ] } private func readExoMetadata(bundle: Bundle = .main) -> (version: String?, commit: String?) { let info = bundle.infoDictionary ?? [:] let tag = info["EXOBuildTag"] as? String let short = info["CFBundleShortVersionString"] as? String let version = [tag, short] .compactMap { $0?.trimmingCharacters(in: .whitespacesAndNewlines) } .first { !$0.isEmpty } let commit = (info["EXOBuildCommit"] as? String)? .trimmingCharacters(in: .whitespacesAndNewlines) let normalizedCommit = (commit?.isEmpty == true) ? nil : commit return (version: version, commit: normalizedCommit) } private func safeRunCommand(_ arguments: [String]) -> String? { let result = runCommand(arguments) guard result.exitCode == 0 else { return nil } let trimmed = result.output.trimmingCharacters(in: .whitespacesAndNewlines) return trimmed.isEmpty ? nil : trimmed } private func extractHardwareUUID(from hardwareProfile: String) -> String? { hardwareProfile .split(separator: "\n") .first { $0.contains("Hardware UUID") }? .split(separator: ":") .dropFirst() .joined(separator: ":") .trimmingCharacters(in: .whitespaces) } private func parseDefaultInterface(from routeOutput: String) -> String? { for line in routeOutput.split(separator: "\n") { let trimmed = line.trimmingCharacters(in: .whitespaces) if trimmed.hasPrefix("interface: ") { return trimmed.replacingOccurrences(of: "interface: ", with: "") } } return nil } private func parseEtherAddress(from ifconfigOutput: String) -> String? { for line in ifconfigOutput.split(separator: "\n") { let trimmed = line.trimmingCharacters(in: .whitespaces) if trimmed.hasPrefix("ether ") { return trimmed.replacingOccurrences(of: "ether ", with: "") } } return nil } private func parseWifiSSID(from airportOutput: String) -> String? { for line in airportOutput.split(separator: "\n") { let trimmed = line.trimmingCharacters(in: .whitespaces) if trimmed.hasPrefix("SSID:") { return trimmed.replacingOccurrences(of: "SSID:", with: "").trimmingCharacters( in: .whitespaces) } } return nil } private func runCommand(_ arguments: [String]) -> CommandResult { let process = Process() process.launchPath = arguments.first process.arguments = Array(arguments.dropFirst()) let stdout = Pipe() let stderr = Pipe() process.standardOutput = stdout process.standardError = stderr do { try process.run() } catch { return CommandResult(exitCode: -1, output: "", error: error.localizedDescription) } process.waitUntilExit() let outputData = stdout.fileHandleForReading.readDataToEndOfFile() let errorData = stderr.fileHandleForReading.readDataToEndOfFile() return CommandResult( exitCode: process.terminationStatus, output: String(decoding: outputData, as: UTF8.self), error: String(decoding: errorData, as: UTF8.self) ) } } private struct DebugInfo { let thunderboltBridgeDisabled: Bool? let interfaces: [InterfaceStatus] let rdma: RDMADebugInfo struct InterfaceStatus { let name: String let ip: String? func toDictionary() -> [String: Any] { [ "name": name, "ip": ip as Any, ] } } struct RDMADebugInfo { let rdmaCtlStatus: String? let ibvDevices: String? let ibvDevinfo: String? func toDictionary() -> [String: Any] { [ "rdma_ctl_status": rdmaCtlStatus as Any, "ibv_devices": ibvDevices as Any, "ibv_devinfo": ibvDevinfo as Any, ] } } func toDictionary() -> [String: Any] { [ "thunderbolt_bridge_disabled": thunderboltBridgeDisabled as Any, "interfaces": interfaces.map { $0.toDictionary() }, "rdma": rdma.toDictionary(), ] } } private struct CommandResult { let exitCode: Int32 let output: String let error: String } ================================================ FILE: app/EXO/EXO/Services/ClusterStateService.swift ================================================ import Combine import Foundation @MainActor final class ClusterStateService: ObservableObject { @Published private(set) var latestSnapshot: ClusterState? @Published private(set) var lastError: String? @Published private(set) var lastActionMessage: String? @Published private(set) var modelOptions: [ModelOption] = [] @Published private(set) var localNodeId: String? private var timer: Timer? private let decoder: JSONDecoder private let session: URLSession private let baseURL: URL private let endpoint: URL init( baseURL: URL = URL(string: "http://127.0.0.1:52415")!, session: URLSession = .shared ) { self.baseURL = baseURL self.endpoint = baseURL.appendingPathComponent("state") self.session = session let decoder = JSONDecoder() decoder.keyDecodingStrategy = .convertFromSnakeCase self.decoder = decoder } func startPolling(interval: TimeInterval = 0.5) { stopPolling() Task { await fetchLocalNodeId() await fetchModels() await fetchSnapshot() } timer = Timer.scheduledTimer(withTimeInterval: interval, repeats: true) { [weak self] _ in Task { await self?.fetchSnapshot() } } } func stopPolling() { timer?.invalidate() timer = nil } func resetTransientState() { latestSnapshot = nil lastError = nil lastActionMessage = nil localNodeId = nil } private func fetchLocalNodeId() async { do { let url = baseURL.appendingPathComponent("node_id") var request = URLRequest(url: url) request.cachePolicy = .reloadIgnoringLocalCacheData let (data, response) = try await session.data(for: request) guard let httpResponse = response as? HTTPURLResponse, (200..<300).contains(httpResponse.statusCode) else { return } if let nodeId = try? decoder.decode(String.self, from: data) { localNodeId = nodeId } } catch { // Silently ignore - localNodeId will remain nil and retry on next poll } } private func fetchSnapshot() async { // Retry fetching local node ID if not yet set if localNodeId == nil { await fetchLocalNodeId() } do { var request = URLRequest(url: endpoint) request.cachePolicy = .reloadIgnoringLocalCacheData let (data, response) = try await session.data(for: request) guard let httpResponse = response as? HTTPURLResponse else { throw URLError(.badServerResponse) } guard (200..<300).contains(httpResponse.statusCode) else { throw URLError(.badServerResponse) } let snapshot = try decoder.decode(ClusterState.self, from: data) latestSnapshot = snapshot if modelOptions.isEmpty { Task { await fetchModels() } } lastError = nil } catch { lastError = error.localizedDescription } } func deleteInstance(_ id: String) async { do { var request = URLRequest(url: baseURL.appendingPathComponent("instance/\(id)")) request.httpMethod = "DELETE" request.setValue("application/json", forHTTPHeaderField: "Accept") let (_, response) = try await session.data(for: request) guard let httpResponse = response as? HTTPURLResponse else { throw URLError(.badServerResponse) } guard (200..<300).contains(httpResponse.statusCode) else { throw URLError(.badServerResponse) } lastActionMessage = "Instance deleted" await fetchSnapshot() } catch { lastError = "Failed to delete instance: \(error.localizedDescription)" } } func launchInstance(modelId: String, sharding: String, instanceMeta: String, minNodes: Int) async { do { var request = URLRequest(url: baseURL.appendingPathComponent("instance")) request.httpMethod = "POST" request.setValue("application/json", forHTTPHeaderField: "Content-Type") let payload: [String: Any] = [ "model_id": modelId, "sharding": sharding, "instance_meta": instanceMeta, "min_nodes": minNodes, ] request.httpBody = try JSONSerialization.data(withJSONObject: payload, options: []) let (_, response) = try await session.data(for: request) guard let httpResponse = response as? HTTPURLResponse else { throw URLError(.badServerResponse) } guard (200..<300).contains(httpResponse.statusCode) else { throw URLError(.badServerResponse) } lastActionMessage = "Instance launched" await fetchSnapshot() } catch { lastError = "Failed to launch instance: \(error.localizedDescription)" } } func fetchModels() async { do { let url = baseURL.appendingPathComponent("models") let (data, response) = try await session.data(from: url) guard let httpResponse = response as? HTTPURLResponse, (200..<300).contains(httpResponse.statusCode) else { throw URLError(.badServerResponse) } let list = try decoder.decode(ModelListResponse.self, from: data) modelOptions = list.data.map { ModelOption(id: $0.id, displayName: $0.name ?? $0.id) } } catch { lastError = "Failed to load models: \(error.localizedDescription)" } } } struct ModelOption: Identifiable { let id: String let displayName: String } struct ModelListResponse: Decodable { let data: [ModelListModel] } struct ModelListModel: Decodable { let id: String let name: String? } ================================================ FILE: app/EXO/EXO/Services/LocalNetworkChecker.swift ================================================ import Foundation import Network import os.log /// Checks if the app's local network permission is actually functional. /// /// macOS local network permission can appear enabled in System Preferences but not /// actually work after a restart. This service uses NWConnection to mDNS multicast /// to verify actual connectivity. @MainActor final class LocalNetworkChecker: ObservableObject { enum Status: Equatable { case unknown case checking case working case notWorking(reason: String) var isHealthy: Bool { if case .working = self { return true } return false } var displayText: String { switch self { case .unknown: return "Unknown" case .checking: return "Checking..." case .working: return "Working" case .notWorking(let reason): return reason } } } private static let logger = Logger(subsystem: "io.exo.EXO", category: "LocalNetworkChecker") private static let hasCompletedInitialCheckKey = "LocalNetworkChecker.hasCompletedInitialCheck" @Published private(set) var status: Status = .unknown private var connection: NWConnection? private var checkTask: Task? private var periodicTask: Task? /// Whether we've completed at least one check (stored in UserDefaults) private var hasCompletedInitialCheck: Bool { get { UserDefaults.standard.bool(forKey: Self.hasCompletedInitialCheckKey) } set { UserDefaults.standard.set(newValue, forKey: Self.hasCompletedInitialCheckKey) } } /// Checks if local network access is working (one-time check). func check() { performCheck() } /// Starts periodic checking of local network access. /// Re-checks every `interval` seconds so the warning disappears when user grants permission. func startPeriodicChecking(interval: TimeInterval = 10) { stopPeriodicChecking() // Do an immediate check first performCheck() // Then schedule periodic checks periodicTask = Task { [weak self] in while !Task.isCancelled { try? await Task.sleep(nanoseconds: UInt64(interval * 1_000_000_000)) guard !Task.isCancelled else { break } self?.performCheck() } } } /// Stops periodic checking. func stopPeriodicChecking() { periodicTask?.cancel() periodicTask = nil } private func performCheck() { checkTask?.cancel() // Only show "checking" status on first check to avoid UI flicker if status == .unknown { status = .checking } // Use longer timeout on first launch to allow time for permission prompt let isFirstCheck = !hasCompletedInitialCheck let timeout: UInt64 = isFirstCheck ? 30_000_000_000 : 3_000_000_000 checkTask = Task { [weak self] in guard let self else { return } Self.logger.debug("Checking local network connectivity (first check: \(isFirstCheck))") let result = await self.checkConnectivity(timeout: timeout) self.status = result self.hasCompletedInitialCheck = true // Only log on state changes or first check to reduce noise if isFirstCheck || result != self.status { Self.logger.info("Local network check: \(result.displayText)") } } } /// Checks connectivity using NWConnection to mDNS multicast. /// The connection attempt triggers the permission prompt if not yet shown. private func checkConnectivity(timeout: UInt64) async -> Status { connection?.cancel() connection = nil // mDNS multicast address - same as libp2p uses for peer discovery let host = NWEndpoint.Host("224.0.0.251") let port = NWEndpoint.Port(integerLiteral: 5353) let params = NWParameters.udp params.allowLocalEndpointReuse = true let conn = NWConnection(host: host, port: port, using: params) connection = conn return await withCheckedContinuation { continuation in var hasResumed = false let lock = NSLock() let resumeOnce: (Status) -> Void = { status in lock.lock() defer { lock.unlock() } guard !hasResumed else { return } hasResumed = true continuation.resume(returning: status) } conn.stateUpdateHandler = { state in switch state { case .ready: resumeOnce(.working) case .waiting(let error): let errorStr = "\(error)" if errorStr.contains("54") || errorStr.contains("ECONNRESET") { resumeOnce(.notWorking(reason: "Connection blocked")) } // Otherwise keep waiting - might be showing permission prompt case .failed(let error): let errorStr = "\(error)" if errorStr.contains("65") || errorStr.contains("EHOSTUNREACH") || errorStr.contains("permission") || errorStr.contains("denied") { resumeOnce(.notWorking(reason: "Permission denied")) } else { resumeOnce(.notWorking(reason: "Failed: \(error.localizedDescription)")) } case .cancelled, .setup, .preparing: break @unknown default: break } } conn.start(queue: .main) Task { try? await Task.sleep(nanoseconds: timeout) let state = conn.state switch state { case .ready: resumeOnce(.working) case .waiting, .preparing, .setup: resumeOnce(.notWorking(reason: "Timeout (may be blocked)")) default: resumeOnce(.notWorking(reason: "Timeout")) } } } } func stop() { stopPeriodicChecking() checkTask?.cancel() checkTask = nil connection?.cancel() connection = nil } } ================================================ FILE: app/EXO/EXO/Services/NetworkSetupHelper.swift ================================================ import AppKit import Foundation import os.log enum NetworkSetupHelper { private static let logger = Logger(subsystem: "io.exo.EXO", category: "NetworkSetup") private static let daemonLabel = "io.exo.networksetup" private static let scriptDestination = "/Library/Application Support/EXO/disable_bridge.sh" // Legacy script path from older versions private static let legacyScriptDestination = "/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh" private static let plistDestination = "/Library/LaunchDaemons/io.exo.networksetup.plist" private static let requiredStartInterval: Int = 1786 private static let setupScript = """ #!/usr/bin/env bash set -euo pipefail # Wait for macOS to finish network setup after boot sleep 20 PREFS="/Library/Preferences/SystemConfiguration/preferences.plist" # Remove bridge0 interface ifconfig bridge0 &>/dev/null && { ifconfig bridge0 | grep -q 'member' && { ifconfig bridge0 | awk '/member/ {print $2}' | xargs -n1 ifconfig bridge0 deletem 2>/dev/null || true } ifconfig bridge0 destroy 2>/dev/null || true } # Remove Thunderbolt Bridge from VirtualNetworkInterfaces in preferences.plist /usr/libexec/PlistBuddy -c "Delete :VirtualNetworkInterfaces:Bridge:bridge0" "$PREFS" 2>/dev/null || true networksetup -listlocations | grep -q exo || { networksetup -createlocation exo } networksetup -switchtolocation exo networksetup -listallhardwareports \\ | awk -F': ' '/Hardware Port: / {print $2}' \\ | while IFS=":" read -r name; do case "$name" in "Ethernet Adapter"*) ;; "Thunderbolt Bridge") ;; "Thunderbolt "*) networksetup -listallnetworkservices \\ | grep -q "EXO $name" \\ || networksetup -createnetworkservice "EXO $name" "$name" 2>/dev/null \\ || continue networksetup -setdhcp "EXO $name" ;; *) networksetup -listallnetworkservices \\ | grep -q "$name" \\ || networksetup -createnetworkservice "$name" "$name" 2>/dev/null \\ || continue ;; esac done networksetup -listnetworkservices | grep -q "Thunderbolt Bridge" && { networksetup -setnetworkserviceenabled "Thunderbolt Bridge" off } || true """ /// Prompts user and installs the LaunchDaemon if not already installed. /// Shows an alert explaining what will be installed before requesting admin privileges. static func promptAndInstallIfNeeded() { // Use .utility priority to match NSAppleScript's internal QoS and avoid priority inversion Task.detached(priority: .utility) { // If already correctly installed, skip if daemonAlreadyInstalled() { return } // Show alert on main thread let shouldInstall = await MainActor.run { let alert = NSAlert() alert.messageText = "EXO Network Configuration" alert.informativeText = "EXO needs to install a system service to configure local networking. This will disable Thunderbolt Bridge (preventing packet storms) and install a Network Location.\n\nYou will be prompted for your password." alert.alertStyle = .informational alert.addButton(withTitle: "Install") alert.addButton(withTitle: "Not Now") return alert.runModal() == .alertFirstButtonReturn } guard shouldInstall else { logger.info("User deferred network setup daemon installation") return } do { try installLaunchDaemon() logger.info("Network setup launch daemon installed and started") } catch { logger.error( "Network setup launch daemon failed: \(error.localizedDescription, privacy: .public)" ) } } } /// Removes all EXO network setup components from the system. /// This includes the LaunchDaemon, scripts, logs, and network location. /// Requires admin privileges. static func uninstall() throws { let uninstallScript = makeUninstallScript() try runShellAsAdmin(uninstallScript) logger.info("EXO network setup components removed successfully") } /// Checks if there are any EXO network components installed that need cleanup static func hasInstalledComponents() -> Bool { let manager = FileManager.default let scriptExists = manager.fileExists(atPath: scriptDestination) let legacyScriptExists = manager.fileExists(atPath: legacyScriptDestination) let plistExists = manager.fileExists(atPath: plistDestination) return scriptExists || legacyScriptExists || plistExists } private static func daemonAlreadyInstalled() -> Bool { let manager = FileManager.default let scriptExists = manager.fileExists(atPath: scriptDestination) let plistExists = manager.fileExists(atPath: plistDestination) guard scriptExists, plistExists else { return false } guard let installedScript = try? String(contentsOfFile: scriptDestination, encoding: .utf8), installedScript.trimmingCharacters(in: .whitespacesAndNewlines) == setupScript.trimmingCharacters(in: .whitespacesAndNewlines) else { return false } guard let data = try? Data(contentsOf: URL(fileURLWithPath: plistDestination)), let plist = try? PropertyListSerialization.propertyList( from: data, options: [], format: nil) as? [String: Any] else { return false } guard let interval = plist["StartInterval"] as? Int, interval == requiredStartInterval else { return false } if let programArgs = plist["ProgramArguments"] as? [String], programArgs.contains(scriptDestination) == false { return false } return true } private static func installLaunchDaemon() throws { let installerScript = makeInstallerScript() try runShellAsAdmin(installerScript) } private static func makeInstallerScript() -> String { """ set -euo pipefail LABEL="\(daemonLabel)" SCRIPT_DEST="\(scriptDestination)" LEGACY_SCRIPT_DEST="\(legacyScriptDestination)" PLIST_DEST="\(plistDestination)" LOG_OUT="/var/log/\(daemonLabel).log" LOG_ERR="/var/log/\(daemonLabel).err.log" # First, completely remove any existing installation launchctl bootout system/"$LABEL" 2>/dev/null || true rm -f "$PLIST_DEST" rm -f "$SCRIPT_DEST" rm -f "$LEGACY_SCRIPT_DEST" rm -f "$LOG_OUT" "$LOG_ERR" # Install fresh mkdir -p "$(dirname "$SCRIPT_DEST")" cat > "$SCRIPT_DEST" <<'EOF_SCRIPT' \(setupScript) EOF_SCRIPT chmod 755 "$SCRIPT_DEST" cat > "$PLIST_DEST" <<'EOF_PLIST' Label \(daemonLabel) ProgramArguments /bin/bash \(scriptDestination) StartInterval \(requiredStartInterval) RunAtLoad StandardOutPath /var/log/\(daemonLabel).log StandardErrorPath /var/log/\(daemonLabel).err.log EOF_PLIST launchctl bootstrap system "$PLIST_DEST" launchctl enable system/"$LABEL" launchctl kickstart -k system/"$LABEL" """ } private static func makeUninstallScript() -> String { """ set -euo pipefail LABEL="\(daemonLabel)" SCRIPT_DEST="\(scriptDestination)" LEGACY_SCRIPT_DEST="\(legacyScriptDestination)" PLIST_DEST="\(plistDestination)" LOG_OUT="/var/log/\(daemonLabel).log" LOG_ERR="/var/log/\(daemonLabel).err.log" # Unload the LaunchDaemon if running launchctl bootout system/"$LABEL" 2>/dev/null || true # Remove LaunchDaemon plist rm -f "$PLIST_DEST" # Remove the script (current and legacy paths) and parent directory if empty rm -f "$SCRIPT_DEST" rm -f "$LEGACY_SCRIPT_DEST" rmdir "$(dirname "$SCRIPT_DEST")" 2>/dev/null || true # Remove log files rm -f "$LOG_OUT" "$LOG_ERR" # Switch back to Automatic network location networksetup -switchtolocation Automatic >/dev/null 2>&1 || true # Delete the exo network location if it exists networksetup -listlocations 2>/dev/null | grep -q '^exo$' && { networksetup -deletelocation exo >/dev/null 2>&1 || true } || true # Re-enable any Thunderbolt Bridge service if it exists # We find it dynamically by looking for bridges containing Thunderbolt interfaces find_and_enable_thunderbolt_bridge() { # Get Thunderbolt interface devices from hardware ports tb_devices=$(networksetup -listallhardwareports 2>/dev/null | awk ' /^Hardware Port:/ { port = tolower(substr($0, 16)) } /^Device:/ { if (port ~ /thunderbolt/) print substr($0, 9) } ') || true [ -z "$tb_devices" ] && return 0 # For each bridge device, check if it contains Thunderbolt interfaces for bridge in bridge0 bridge1 bridge2; do members=$(ifconfig "$bridge" 2>/dev/null | awk '/member:/ {print $2}') || true [ -z "$members" ] && continue for tb_dev in $tb_devices; do if echo "$members" | grep -qx "$tb_dev"; then # Find the service name for this bridge device service_name=$(networksetup -listnetworkserviceorder 2>/dev/null | awk -v dev="$bridge" ' /^\\([0-9*]/ { gsub(/^\\([0-9*]+\\) /, ""); svc = $0 } /Device:/ && $0 ~ dev { print svc; exit } ') || true if [ -n "$service_name" ]; then networksetup -setnetworkserviceenabled "$service_name" on 2>/dev/null || true return 0 fi fi done done return 0 } find_and_enable_thunderbolt_bridge || true echo "EXO network components removed successfully" """ } /// Direct install without GUI (requires root). /// Returns true on success, false on failure. static func installDirectly() -> Bool { let script = makeInstallerScript() return runShellDirectly(script) } /// Direct uninstall without GUI (requires root). /// Returns true on success, false on failure. static func uninstallDirectly() -> Bool { let script = makeUninstallScript() return runShellDirectly(script) } /// Run a shell script directly via Process (no AppleScript, requires root). /// Returns true on success, false on failure. private static func runShellDirectly(_ script: String) -> Bool { let process = Process() process.executableURL = URL(fileURLWithPath: "/bin/bash") process.arguments = ["-c", script] let outputPipe = Pipe() let errorPipe = Pipe() process.standardOutput = outputPipe process.standardError = errorPipe do { try process.run() process.waitUntilExit() let outputData = outputPipe.fileHandleForReading.readDataToEndOfFile() let errorData = errorPipe.fileHandleForReading.readDataToEndOfFile() if let output = String(data: outputData, encoding: .utf8), !output.isEmpty { print(output) } if let errorOutput = String(data: errorData, encoding: .utf8), !errorOutput.isEmpty { fputs(errorOutput, stderr) } if process.terminationStatus == 0 { logger.info("Shell script completed successfully") return true } else { logger.error("Shell script failed with exit code \(process.terminationStatus)") return false } } catch { logger.error( "Failed to run shell script: \(error.localizedDescription, privacy: .public)") fputs("Error: \(error.localizedDescription)\n", stderr) return false } } private static func runShellAsAdmin(_ script: String) throws { let escapedScript = script .replacingOccurrences(of: "\\", with: "\\\\") .replacingOccurrences(of: "\"", with: "\\\"") let appleScriptSource = """ do shell script "\(escapedScript)" with administrator privileges """ guard let appleScript = NSAppleScript(source: appleScriptSource) else { throw NetworkSetupError.scriptCreationFailed } var errorInfo: NSDictionary? appleScript.executeAndReturnError(&errorInfo) if let errorInfo { let message = errorInfo[NSAppleScript.errorMessage] as? String ?? "Unknown error" throw NetworkSetupError.executionFailed(message) } } } enum NetworkSetupError: LocalizedError { case scriptCreationFailed case executionFailed(String) var errorDescription: String? { switch self { case .scriptCreationFailed: return "Failed to create AppleScript for network setup" case .executionFailed(let message): return "Network setup script failed: \(message)" } } } ================================================ FILE: app/EXO/EXO/Services/NetworkStatusService.swift ================================================ import AppKit import Foundation @MainActor final class NetworkStatusService: ObservableObject { @Published private(set) var status: NetworkStatus = .empty private var timer: Timer? func refresh() async { let fetched = await Task.detached(priority: .background) { NetworkStatusFetcher().fetch() }.value status = fetched } func startPolling(interval: TimeInterval = 30) { timer?.invalidate() timer = Timer.scheduledTimer(withTimeInterval: interval, repeats: true) { [weak self] _ in guard let self else { return } Task { await self.refresh() } } if let timer { RunLoop.main.add(timer, forMode: .common) } Task { await refresh() } } func stopPolling() { timer?.invalidate() timer = nil } } struct NetworkStatus: Equatable { let thunderboltBridgeState: ThunderboltState? let bridgeInactive: Bool? let interfaceStatuses: [InterfaceIpStatus] let localRdmaDevices: [String] let localRdmaActivePorts: [RDMAPort] static let empty = NetworkStatus( thunderboltBridgeState: nil, bridgeInactive: nil, interfaceStatuses: [], localRdmaDevices: [], localRdmaActivePorts: [] ) } struct RDMAPort: Equatable { let device: String let port: String let state: String } struct InterfaceIpStatus: Equatable { let interfaceName: String let ipAddress: String? } enum ThunderboltState: Equatable { case enabled case disabled case deleted } private struct NetworkStatusFetcher { func fetch() -> NetworkStatus { NetworkStatus( thunderboltBridgeState: readThunderboltBridgeState(), bridgeInactive: readBridgeInactive(), interfaceStatuses: readInterfaceStatuses(), localRdmaDevices: readRDMADevices(), localRdmaActivePorts: readRDMAActivePorts() ) } private func readRDMADevices() -> [String] { let result = runCommand(["ibv_devices"]) guard result.exitCode == 0 else { return [] } var devices: [String] = [] for line in result.output.split(separator: "\n") { let trimmed = line.trimmingCharacters(in: .whitespaces) if trimmed.hasPrefix("---") || trimmed.lowercased().hasPrefix("device") || trimmed.isEmpty { continue } let parts = trimmed.split(separator: " ", maxSplits: 1) if let deviceName = parts.first { devices.append(String(deviceName)) } } return devices } private func readRDMAActivePorts() -> [RDMAPort] { let result = runCommand(["ibv_devinfo"]) guard result.exitCode == 0 else { return [] } var ports: [RDMAPort] = [] var currentDevice: String? var currentPort: String? for line in result.output.split(separator: "\n") { let trimmed = line.trimmingCharacters(in: .whitespaces) if trimmed.hasPrefix("hca_id:") { currentDevice = trimmed.replacingOccurrences(of: "hca_id:", with: "") .trimmingCharacters(in: .whitespaces) } else if trimmed.hasPrefix("port:") { currentPort = trimmed.replacingOccurrences(of: "port:", with: "") .trimmingCharacters(in: .whitespaces) } else if trimmed.hasPrefix("state:") { let state = trimmed.replacingOccurrences(of: "state:", with: "").trimmingCharacters( in: .whitespaces) if let device = currentDevice, let port = currentPort { if state.lowercased().contains("active") { ports.append(RDMAPort(device: device, port: port, state: state)) } } } } return ports } private func readThunderboltBridgeState() -> ThunderboltState? { // Dynamically find the Thunderbolt Bridge service (don't assume the name) guard let serviceName = ThunderboltBridgeDetector.findThunderboltBridgeServiceName() else { // No bridge containing Thunderbolt interfaces exists return .deleted } guard let isEnabled = ThunderboltBridgeDetector.isServiceEnabled(serviceName: serviceName) else { return nil } return isEnabled ? .enabled : .disabled } private func readBridgeInactive() -> Bool? { let result = runCommand(["ifconfig", "bridge0"]) guard result.exitCode == 0 else { return nil } guard let statusLine = result.output .components(separatedBy: .newlines) .first(where: { $0.contains("status:") })? .lowercased() else { return nil } if statusLine.contains("inactive") { return true } if statusLine.contains("active") { return false } return nil } private func readInterfaceStatuses() -> [InterfaceIpStatus] { (0...7).map { "en\($0)" }.map(readInterfaceStatus) } private func readInterfaceStatus(for interface: String) -> InterfaceIpStatus { let result = runCommand(["ifconfig", interface]) guard result.exitCode == 0 else { return InterfaceIpStatus( interfaceName: interface, ipAddress: nil ) } let output = result.output let ip = firstInet(from: output) return InterfaceIpStatus( interfaceName: interface, ipAddress: ip ) } private func firstInet(from ifconfigOutput: String) -> String? { for line in ifconfigOutput.split(separator: "\n") { let trimmed = line.trimmingCharacters(in: .whitespaces) guard trimmed.hasPrefix("inet ") else { continue } let parts = trimmed.split(separator: " ") if parts.count >= 2 { let candidate = String(parts[1]) if candidate != "127.0.0.1" { return candidate } } } return nil } private struct CommandResult { let exitCode: Int32 let output: String let error: String } private func runCommand(_ arguments: [String]) -> CommandResult { let process = Process() process.launchPath = "/usr/bin/env" process.arguments = arguments let stdout = Pipe() let stderr = Pipe() process.standardOutput = stdout process.standardError = stderr do { try process.run() } catch { return CommandResult(exitCode: -1, output: "", error: error.localizedDescription) } process.waitUntilExit() let outputData = stdout.fileHandleForReading.readDataToEndOfFile() let errorData = stderr.fileHandleForReading.readDataToEndOfFile() return CommandResult( exitCode: process.terminationStatus, output: String(decoding: outputData, as: UTF8.self), error: String(decoding: errorData, as: UTF8.self) ) } } ================================================ FILE: app/EXO/EXO/Services/ThunderboltBridgeDetector.swift ================================================ import Foundation import os.log /// Utility for dynamically detecting Thunderbolt Bridge network services. /// This mirrors the Python logic in info_gatherer.py - we never assume the service /// is named "Thunderbolt Bridge", instead we find bridges containing Thunderbolt interfaces. enum ThunderboltBridgeDetector { private static let logger = Logger( subsystem: "io.exo.EXO", category: "ThunderboltBridgeDetector") struct CommandResult { let exitCode: Int32 let output: String let error: String } /// Find the network service name of a bridge containing Thunderbolt interfaces. /// Returns nil if no such bridge exists. static func findThunderboltBridgeServiceName() -> String? { // 1. Get all Thunderbolt interface devices (e.g., en2, en3) guard let thunderboltDevices = getThunderboltDevices(), !thunderboltDevices.isEmpty else { logger.debug("No Thunderbolt devices found") return nil } logger.debug("Found Thunderbolt devices: \(thunderboltDevices.joined(separator: ", "))") // 2. Get bridge services from network service order guard let bridgeServices = getBridgeServices(), !bridgeServices.isEmpty else { logger.debug("No bridge services found") return nil } logger.debug("Found bridge services: \(bridgeServices.keys.joined(separator: ", "))") // 3. Find a bridge that contains Thunderbolt interfaces for (bridgeDevice, serviceName) in bridgeServices { let members = getBridgeMembers(bridgeDevice: bridgeDevice) logger.debug( "Bridge \(bridgeDevice) (\(serviceName)) has members: \(members.joined(separator: ", "))" ) // Check if any Thunderbolt device is a member of this bridge if !members.isDisjoint(with: thunderboltDevices) { logger.info( "Found Thunderbolt Bridge service: '\(serviceName)' (device: \(bridgeDevice))") return serviceName } } logger.debug("No bridge found containing Thunderbolt interfaces") return nil } /// Get Thunderbolt interface device names (e.g., en2, en3) from hardware ports. private static func getThunderboltDevices() -> Set? { let result = runCommand(["networksetup", "-listallhardwareports"]) guard result.exitCode == 0 else { logger.warning("networksetup -listallhardwareports failed: \(result.error)") return nil } var thunderboltDevices: Set = [] var currentPort: String? for line in result.output.components(separatedBy: .newlines) { let trimmed = line.trimmingCharacters(in: .whitespaces) if trimmed.hasPrefix("Hardware Port:") { currentPort = String(trimmed.dropFirst("Hardware Port:".count)).trimmingCharacters( in: .whitespaces) } else if trimmed.hasPrefix("Device:"), let port = currentPort { let device = String(trimmed.dropFirst("Device:".count)).trimmingCharacters( in: .whitespaces) if port.lowercased().contains("thunderbolt") { thunderboltDevices.insert(device) } currentPort = nil } } return thunderboltDevices } /// Get mapping of bridge device -> service name from network service order. private static func getBridgeServices() -> [String: String]? { let result = runCommand(["networksetup", "-listnetworkserviceorder"]) guard result.exitCode == 0 else { logger.warning("networksetup -listnetworkserviceorder failed: \(result.error)") return nil } // Parse service order to find bridge devices and their service names // Format: "(1) Service Name\n(Hardware Port: ..., Device: bridge0)\n" var bridgeServices: [String: String] = [:] var currentService: String? for line in result.output.components(separatedBy: .newlines) { let trimmed = line.trimmingCharacters(in: .whitespaces) // Match "(N) Service Name" or "(*) Service Name" (disabled) // but NOT "(Hardware Port: ...)" lines if trimmed.hasPrefix("("), trimmed.contains(")"), !trimmed.hasPrefix("(Hardware Port:") { if let parenEnd = trimmed.firstIndex(of: ")") { let afterParen = trimmed.index(after: parenEnd) if afterParen < trimmed.endIndex { currentService = String(trimmed[afterParen...]) .trimmingCharacters(in: .whitespaces) } } } // Match "(Hardware Port: ..., Device: bridgeX)" else if let service = currentService, trimmed.contains("Device: bridge") { // Extract device name from "..., Device: bridge0)" if let deviceRange = trimmed.range(of: "Device: ") { let afterDevice = trimmed[deviceRange.upperBound...] if let parenIndex = afterDevice.firstIndex(of: ")") { let device = String(afterDevice[.. Set { let result = runCommand(["ifconfig", bridgeDevice]) guard result.exitCode == 0 else { logger.debug("ifconfig \(bridgeDevice) failed") return [] } var members: Set = [] for line in result.output.components(separatedBy: .newlines) { let trimmed = line.trimmingCharacters(in: .whitespaces) if trimmed.hasPrefix("member:") { let parts = trimmed.split(separator: " ") if parts.count > 1 { members.insert(String(parts[1])) } } } return members } /// Check if a network service is enabled. static func isServiceEnabled(serviceName: String) -> Bool? { let result = runCommand(["networksetup", "-getnetworkserviceenabled", serviceName]) guard result.exitCode == 0 else { logger.warning("Failed to check if '\(serviceName)' is enabled: \(result.error)") return nil } let output = result.output.lowercased().trimmingCharacters(in: .whitespacesAndNewlines) if output.contains("enabled") { return true } if output.contains("disabled") { return false } return nil } private static func runCommand(_ arguments: [String]) -> CommandResult { let process = Process() process.launchPath = "/usr/bin/env" process.arguments = arguments let stdout = Pipe() let stderr = Pipe() process.standardOutput = stdout process.standardError = stderr do { try process.run() } catch { return CommandResult(exitCode: -1, output: "", error: error.localizedDescription) } process.waitUntilExit() let outputData = stdout.fileHandleForReading.readDataToEndOfFile() let errorData = stderr.fileHandleForReading.readDataToEndOfFile() return CommandResult( exitCode: process.terminationStatus, output: String(decoding: outputData, as: UTF8.self), error: String(decoding: errorData, as: UTF8.self) ) } } ================================================ FILE: app/EXO/EXO/Services/ThunderboltBridgeService.swift ================================================ import AppKit import Combine import Foundation import Security import SystemConfiguration import os.log @MainActor final class ThunderboltBridgeService: ObservableObject { private static let logger = Logger(subsystem: "io.exo.EXO", category: "ThunderboltBridge") @Published private(set) var detectedCycle: [String]? @Published private(set) var hasPromptedForCurrentCycle = false @Published private(set) var lastError: String? private weak var clusterStateService: ClusterStateService? private var cancellables = Set() private var previousCycleSignature: String? init(clusterStateService: ClusterStateService) { self.clusterStateService = clusterStateService setupObserver() } private func setupObserver() { guard let service = clusterStateService else { return } service.$latestSnapshot .compactMap { $0 } .sink { [weak self] snapshot in self?.checkForCycles(snapshot: snapshot) } .store(in: &cancellables) } private func checkForCycles(snapshot: ClusterState) { let cycles = snapshot.thunderboltBridgeCycles // Only consider cycles with more than 2 nodes guard let firstCycle = cycles.first, firstCycle.count > 2 else { // No problematic cycles detected, reset state if detectedCycle != nil { detectedCycle = nil previousCycleSignature = nil hasPromptedForCurrentCycle = false } return } // Create a signature for this cycle to detect if it changed let cycleSignature = firstCycle.sorted().joined(separator: ",") // If this is a new/different cycle, reset the prompt state if cycleSignature != previousCycleSignature { previousCycleSignature = cycleSignature hasPromptedForCurrentCycle = false } detectedCycle = firstCycle // Only prompt once per cycle if !hasPromptedForCurrentCycle { showDisableBridgePrompt(nodeIds: firstCycle) } } private func showDisableBridgePrompt(nodeIds: [String]) { hasPromptedForCurrentCycle = true // Get friendly names for the nodes if available let nodeNames = nodeIds.map { nodeId -> String in if let snapshot = clusterStateService?.latestSnapshot, let profile = snapshot.nodeProfiles[nodeId], let friendlyName = profile.friendlyName, !friendlyName.isEmpty { return friendlyName } return String(nodeId.prefix(8)) // Use first 8 chars of node ID as fallback } let machineNames = nodeNames.joined(separator: ", ") let alert = NSAlert() alert.messageText = "Thunderbolt Bridge Loop Detected" alert.informativeText = """ A Thunderbolt Bridge loop has been detected between \(nodeNames.count) machines: \(machineNames). This can cause network packet storms and connectivity issues. Would you like to disable Thunderbolt Bridge on this machine to break the loop? """ alert.alertStyle = .warning alert.addButton(withTitle: "Disable Bridge") alert.addButton(withTitle: "Not Now") let response = alert.runModal() if response == .alertFirstButtonReturn { Task { await disableThunderboltBridge() } } } func disableThunderboltBridge() async { Self.logger.info("Attempting to disable Thunderbolt Bridge via SCPreferences") lastError = nil do { try await disableThunderboltBridgeWithSCPreferences() Self.logger.info("Successfully disabled Thunderbolt Bridge") } catch { Self.logger.error( "Failed to disable Thunderbolt Bridge: \(error.localizedDescription, privacy: .public)" ) lastError = error.localizedDescription showErrorAlert(message: error.localizedDescription) } } private func disableThunderboltBridgeWithSCPreferences() async throws { // 1. Create authorization reference var authRef: AuthorizationRef? var status = AuthorizationCreate(nil, nil, [], &authRef) guard status == errAuthorizationSuccess, let authRef = authRef else { throw ThunderboltBridgeError.authorizationFailed } defer { AuthorizationFree(authRef, [.destroyRights]) } // 2. Request specific network configuration rights let rightName = "system.services.systemconfiguration.network" status = rightName.withCString { nameCString in var item = AuthorizationItem( name: nameCString, valueLength: 0, value: nil, flags: 0 ) return withUnsafeMutablePointer(to: &item) { itemPointer in var rights = AuthorizationRights(count: 1, items: itemPointer) return AuthorizationCopyRights( authRef, &rights, nil, [.extendRights, .interactionAllowed], nil ) } } guard status == errAuthorizationSuccess else { if status == errAuthorizationCanceled { throw ThunderboltBridgeError.authorizationCanceled } throw ThunderboltBridgeError.authorizationDenied } // 3. Create SCPreferences with authorization guard let prefs = SCPreferencesCreateWithAuthorization( kCFAllocatorDefault, "EXO" as CFString, nil, authRef ) else { throw ThunderboltBridgeError.preferencesCreationFailed } // 4. Lock, modify, commit guard SCPreferencesLock(prefs, true) else { throw ThunderboltBridgeError.lockFailed } defer { SCPreferencesUnlock(prefs) } // 5. Find the Thunderbolt Bridge service dynamically (don't assume the name) guard let targetServiceName = ThunderboltBridgeDetector.findThunderboltBridgeServiceName() else { throw ThunderboltBridgeError.serviceNotFound } guard let allServices = SCNetworkServiceCopyAll(prefs) as? [SCNetworkService] else { throw ThunderboltBridgeError.servicesNotFound } var found = false for service in allServices { if let name = SCNetworkServiceGetName(service) as String?, name == targetServiceName { guard SCNetworkServiceSetEnabled(service, false) else { throw ThunderboltBridgeError.disableFailed } found = true Self.logger.info( "Found and disabled Thunderbolt Bridge service: '\(targetServiceName)'") break } } if !found { throw ThunderboltBridgeError.serviceNotFound } // 6. Commit and apply guard SCPreferencesCommitChanges(prefs) else { throw ThunderboltBridgeError.commitFailed } guard SCPreferencesApplyChanges(prefs) else { throw ThunderboltBridgeError.applyFailed } } private func showErrorAlert(message: String) { let alert = NSAlert() alert.messageText = "Failed to Disable Thunderbolt Bridge" alert.informativeText = message alert.alertStyle = .critical alert.addButton(withTitle: "OK") alert.runModal() } } enum ThunderboltBridgeError: LocalizedError { case authorizationFailed case authorizationCanceled case authorizationDenied case preferencesCreationFailed case lockFailed case servicesNotFound case serviceNotFound case disableFailed case commitFailed case applyFailed var errorDescription: String? { switch self { case .authorizationFailed: return "Failed to create authorization" case .authorizationCanceled: return "Authorization was canceled by user" case .authorizationDenied: return "Authorization was denied" case .preferencesCreationFailed: return "Failed to access network preferences" case .lockFailed: return "Failed to lock network preferences for modification" case .servicesNotFound: return "Could not retrieve network services" case .serviceNotFound: return "Thunderbolt Bridge service not found" case .disableFailed: return "Failed to disable Thunderbolt Bridge service" case .commitFailed: return "Failed to save network configuration changes" case .applyFailed: return "Failed to apply network configuration changes" } } } ================================================ FILE: app/EXO/EXO/ViewModels/InstanceViewModel.swift ================================================ import Foundation struct DownloadProgressViewModel: Equatable { let downloadedBytes: Int64 let totalBytes: Int64 let speedBytesPerSecond: Double let etaSeconds: Double? let completedFiles: Int let totalFiles: Int var fractionCompleted: Double { guard totalBytes > 0 else { return 0 } return Double(downloadedBytes) / Double(totalBytes) } var percentCompleted: Double { fractionCompleted * 100 } var formattedProgress: String { let downloaded = formatBytes(downloadedBytes) let total = formatBytes(totalBytes) let percent = String(format: "%.1f", percentCompleted) return "\(downloaded)/\(total) (\(percent)%)" } var formattedSpeed: String { "\(formatBytes(Int64(speedBytesPerSecond)))/s" } var formattedETA: String? { guard let eta = etaSeconds, eta > 0 else { return nil } let minutes = Int(eta) / 60 let seconds = Int(eta) % 60 if minutes > 0 { return "ETA \(minutes)m \(seconds)s" } return "ETA \(seconds)s" } private func formatBytes(_ bytes: Int64) -> String { let gb = Double(bytes) / 1_073_741_824.0 let mb = Double(bytes) / 1_048_576.0 if gb >= 1.0 { return String(format: "%.2f GB", gb) } return String(format: "%.0f MB", mb) } } struct InstanceViewModel: Identifiable, Equatable { enum State { case downloading case warmingUp case running case ready case waiting case failed case idle case preparing var label: String { switch self { case .downloading: return "Downloading" case .warmingUp: return "Warming Up" case .running: return "Running" case .ready: return "Ready" case .waiting: return "Waiting" case .failed: return "Failed" case .idle: return "Idle" case .preparing: return "Preparing" } } } let id: String let modelName: String let sharding: String? let nodeNames: [String] let state: State let chatTasks: [InstanceTaskViewModel] let downloadProgress: DownloadProgressViewModel? var nodeSummary: String { guard !nodeNames.isEmpty else { return "0 nodes" } if nodeNames.count == 1 { return nodeNames[0] } if nodeNames.count == 2 { return nodeNames.joined(separator: ", ") } let others = nodeNames.count - 1 return "\(nodeNames.first ?? "") +\(others)" } } extension ClusterState { func instanceViewModels() -> [InstanceViewModel] { let chatTasksByInstance = Dictionary( grouping: tasks.values.filter { $0.kind == .chatCompletion && $0.instanceId != nil }, by: { $0.instanceId! } ) return instances.map { entry in let instance = entry.value let modelName = instance.shardAssignments.modelId let nodeToRunner = instance.shardAssignments.nodeToRunner let nodeIds = Array(nodeToRunner.keys) let runnerIds = Array(nodeToRunner.values) let nodeNames = nodeIds.compactMap { nodeProfiles[$0]?.friendlyName ?? nodeProfiles[$0]?.modelId ?? $0 } let statuses = runnerIds.compactMap { runners[$0]?.status.lowercased() } let downloadProgress = aggregateDownloadProgress(for: nodeIds) let state = InstanceViewModel.State( statuses: statuses, hasActiveDownload: downloadProgress != nil) let chatTasks = (chatTasksByInstance[entry.key] ?? []) .sorted(by: { $0.sortPriority < $1.sortPriority }) .map { InstanceTaskViewModel(task: $0) } return InstanceViewModel( id: entry.key, modelName: modelName, sharding: InstanceViewModel.friendlyShardingName(for: instance.sharding), nodeNames: nodeNames, state: state, chatTasks: chatTasks, downloadProgress: downloadProgress ) } .sorted { $0.modelName < $1.modelName } } private func aggregateDownloadProgress(for nodeIds: [String]) -> DownloadProgressViewModel? { var totalDownloaded: Int64 = 0 var totalSize: Int64 = 0 var totalSpeed: Double = 0 var maxEtaMs: Int64 = 0 var totalCompletedFiles = 0 var totalFileCount = 0 var hasActiveDownload = false for nodeId in nodeIds { guard let nodeDownloads = downloads[nodeId] else { continue } for download in nodeDownloads { guard let progress = download.progress else { continue } hasActiveDownload = true totalDownloaded += progress.downloadedBytes.inBytes totalSize += progress.totalBytes.inBytes totalSpeed += progress.speed ?? 0 if let eta = progress.etaMs { maxEtaMs = max(maxEtaMs, eta) } totalCompletedFiles += progress.completedFiles ?? 0 totalFileCount += progress.totalFiles ?? 0 } } guard hasActiveDownload else { return nil } return DownloadProgressViewModel( downloadedBytes: totalDownloaded, totalBytes: totalSize, speedBytesPerSecond: totalSpeed, etaSeconds: maxEtaMs > 0 ? Double(maxEtaMs) / 1000.0 : nil, completedFiles: totalCompletedFiles, totalFiles: totalFileCount ) } } extension InstanceViewModel.State { fileprivate init(statuses: [String], hasActiveDownload: Bool = false) { if statuses.contains(where: { $0.contains("failed") }) { self = .failed } else if hasActiveDownload || statuses.contains(where: { $0.contains("downloading") }) { self = .downloading } else if statuses.contains(where: { $0.contains("warming") }) { self = .warmingUp } else if statuses.contains(where: { $0.contains("running") }) { self = .running } else if statuses.contains(where: { $0.contains("ready") || $0.contains("loaded") }) { self = .ready } else if statuses.contains(where: { $0.contains("waiting") }) { self = .waiting } else if statuses.isEmpty { self = .idle } else { self = .preparing } } } extension InstanceViewModel { static func friendlyShardingName(for raw: String?) -> String? { guard let raw else { return nil } switch raw.lowercased() { case "mlxringinstance", "mlxring": return "MLX Ring" case "mlxibvinstance", "mlxibv": return "MLX RDMA" default: return raw } } } struct InstanceTaskViewModel: Identifiable, Equatable { enum Kind { case chatCompletion } let id: String let kind: Kind let status: TaskStatus let modelName: String? let promptPreview: String? let errorMessage: String? let subtitle: String? let parameters: TextGenerationTaskParameters? var title: String { switch kind { case .chatCompletion: return "Chat Completion" } } var detailText: String? { if let errorMessage, !errorMessage.isEmpty { return errorMessage } return promptPreview } } extension InstanceTaskViewModel { init(task: ClusterTask) { self.id = task.id self.kind = .chatCompletion self.status = task.status self.modelName = task.modelName self.promptPreview = task.promptPreview self.errorMessage = task.errorMessage self.subtitle = task.modelName self.parameters = task.parameters } } ================================================ FILE: app/EXO/EXO/ViewModels/NodeViewModel.swift ================================================ import Foundation struct NodeViewModel: Identifiable, Equatable { let id: String let friendlyName: String let model: String let usedRamGB: Double let totalRamGB: Double let gpuUsagePercent: Double let cpuUsagePercent: Double let temperatureCelsius: Double let systenPowerWatts: Double var memoryProgress: Double { guard totalRamGB > 0 else { return 0 } return min(max(usedRamGB / totalRamGB, 0), 1) } var memoryLabel: String { String(format: "%.1f / %.1f GB", usedRamGB, totalRamGB) } var temperatureLabel: String { String(format: "%.0f°C", temperatureCelsius) } var powerLabel: String { systenPowerWatts > 0 ? String(format: "%.0fW", systenPowerWatts) : "—" } var cpuUsageLabel: String { String(format: "%.0f%%", cpuUsagePercent) } var gpuUsageLabel: String { String(format: "%.0f%%", gpuUsagePercent) } var deviceIconName: String { let lower = model.lowercased() if lower.contains("studio") { return "macstudio" } if lower.contains("mini") { return "macmini" } return "macbook" } } extension ClusterState { func nodeViewModels() -> [NodeViewModel] { nodeProfiles.map { entry in let profile = entry.value let friendly = profile.friendlyName ?? profile.modelId ?? entry.key let model = profile.modelId ?? "Unknown" let totalBytes = Double(profile.memory?.ramTotal?.inBytes ?? 0) let availableBytes = Double(profile.memory?.ramAvailable?.inBytes ?? 0) let usedBytes = max(totalBytes - availableBytes, 0) return NodeViewModel( id: entry.key, friendlyName: friendly, model: model, usedRamGB: usedBytes / 1_073_741_824.0, totalRamGB: totalBytes / 1_073_741_824.0, gpuUsagePercent: (profile.system?.gpuUsage ?? 0) * 100, cpuUsagePercent: (profile.system?.pcpuUsage ?? 0) * 100, temperatureCelsius: profile.system?.temp ?? 0, systenPowerWatts: profile.system?.sysPower ?? 0 ) } .sorted { $0.friendlyName < $1.friendlyName } } } struct TopologyEdgeViewModel: Hashable { let sourceId: String let targetId: String } struct TopologyViewModel { let nodes: [NodeViewModel] let edges: [TopologyEdgeViewModel] let currentNodeId: String? } extension ClusterState { func topologyViewModel(localNodeId: String?) -> TopologyViewModel? { let topologyNodeIds = Set(topology?.nodes ?? []) let allNodes = nodeViewModels().filter { topologyNodeIds.isEmpty || topologyNodeIds.contains($0.id) } guard !allNodes.isEmpty else { return nil } let nodesById = Dictionary(uniqueKeysWithValues: allNodes.map { ($0.id, $0) }) var orderedNodes: [NodeViewModel] = [] if let topologyNodes = topology?.nodes { for nodeId in topologyNodes { if let viewModel = nodesById[nodeId] { orderedNodes.append(viewModel) } } let seenIds = Set(orderedNodes.map(\.id)) let remaining = allNodes.filter { !seenIds.contains($0.id) } orderedNodes.append(contentsOf: remaining) } else { orderedNodes = allNodes } // Rotate so the local node (from /node_id API) is first if let localId = localNodeId, let index = orderedNodes.firstIndex(where: { $0.id == localId }) { orderedNodes = Array(orderedNodes[index...]) + Array(orderedNodes[..? private static let dashboardURL = "http://localhost:52415/" /// Called when the user completes onboarding (clicks Open Dashboard or dismisses). var onComplete: (() -> Void)? func show() { guard popover == nil else { return } // The status bar button may not exist yet on first launch; retry generously. showWithRetry(attemptsRemaining: 15) } private func showWithRetry(attemptsRemaining: Int) { guard attemptsRemaining > 0 else { // Exhausted retries — fall back to just opening the dashboard directly. openDashboard() onComplete?() return } guard let button = Self.findStatusItemButton() else { DispatchQueue.main.asyncAfter(deadline: .now() + 1.0) { [weak self] in self?.showWithRetry(attemptsRemaining: attemptsRemaining - 1) } return } let pop = NSPopover() pop.behavior = .applicationDefined pop.animates = true pop.contentSize = NSSize(width: 280, height: 120) pop.contentViewController = NSHostingController( rootView: WelcomeCalloutView( countdownDuration: 10, onDismiss: { [weak self] in self?.onComplete?() self?.dismiss() }, onOpen: { [weak self] in self?.openDashboard() self?.onComplete?() self?.dismiss() } ) ) self.popover = pop pop.show(relativeTo: button.bounds, of: button, preferredEdge: .minY) // Auto-open dashboard after 10s then dismiss countdownTask = Task { try? await Task.sleep(nanoseconds: 10_000_000_000) if !Task.isCancelled { openDashboard() onComplete?() dismiss() } } } func dismiss() { countdownTask?.cancel() countdownTask = nil guard let pop = popover else { return } popover = nil pop.performClose(nil) } private func openDashboard() { guard let url = URL(string: Self.dashboardURL) else { return } NSWorkspace.shared.open(url) } /// Finds the NSStatusBarButton created by SwiftUI's MenuBarExtra. /// Walks the view hierarchy to find the actual button rather than the content view. private static func findStatusItemButton() -> NSView? { for window in NSApp.windows { let className = NSStringFromClass(type(of: window)) // Match NSStatusBarWindow or any internal SwiftUI status bar window guard className.contains("StatusBar") || className.contains("MenuBarExtra") else { continue } if let content = window.contentView { if let button = findButton(in: content) { return button } // Fall back to the content view itself if it has a non-zero frame if content.frame.width > 0 { return content } } } return nil } /// Recursively searches the view hierarchy for an NSStatusBarButton. private static func findButton(in view: NSView) -> NSView? { let className = NSStringFromClass(type(of: view)) if className.contains("StatusBarButton") || className.contains("StatusItem") { return view } for subview in view.subviews { if let found = findButton(in: subview) { return found } } return nil } } /// Minimal welcome callout — friendly pointer, not a wall of text. /// Rendered inside the NSPopover which provides its own chrome and arrow. private struct WelcomeCalloutView: View { let countdownDuration: Int let onDismiss: () -> Void let onOpen: () -> Void @State private var countdown: Int @State private var timerTask: Task? init(countdownDuration: Int, onDismiss: @escaping () -> Void, onOpen: @escaping () -> Void) { self.countdownDuration = countdownDuration self.onDismiss = onDismiss self.onOpen = onOpen self._countdown = State(initialValue: countdownDuration) } var body: some View { VStack(alignment: .leading, spacing: 10) { HStack(alignment: .top) { Text("EXO is running") .font(.system(.headline, design: .rounded)) .fontWeight(.semibold) .foregroundColor(.primary) Spacer() Button { onDismiss() } label: { Image(systemName: "xmark.circle.fill") .font(.system(size: 14)) .foregroundStyle(.tertiary) } .buttonStyle(.plain) } Text("Run your first model here:") .font(.system(.subheadline, design: .default)) .foregroundColor(.secondary) HStack { Button { onOpen() } label: { Label("Open Dashboard", systemImage: "arrow.up.right.square") .font(.system(.caption, design: .default)) .fontWeight(.medium) } .buttonStyle(.borderedProminent) .tint(.accentColor) .controlSize(.small) Spacer() if countdown > 0 { Text("Launching in \(countdown) secs...") .font(.system(.caption2, design: .default)) .foregroundColor(.secondary.opacity(0.6)) .monospacedDigit() } } } .padding(14) .onAppear { startCountdown() } .onDisappear { timerTask?.cancel() timerTask = nil } } private func startCountdown() { timerTask = Task { while countdown > 0 { try? await Task.sleep(nanoseconds: 1_000_000_000) if !Task.isCancelled { countdown -= 1 } } } } } ================================================ FILE: app/EXO/EXO/Views/InstanceRowView.swift ================================================ import SwiftUI struct InstanceRowView: View { let instance: InstanceViewModel @State private var animatedTaskIDs: Set = [] @State private var infoTask: InstanceTaskViewModel? @State private var showChatTasks = true var body: some View { VStack(alignment: .leading, spacing: 6) { HStack(spacing: 8) { VStack(alignment: .leading, spacing: 2) { Text(instance.modelName) .font(.subheadline) Text(instance.nodeSummary) .font(.caption) .foregroundColor(.secondary) } Spacer() if let progress = instance.downloadProgress { downloadStatusView(progress: progress) } else { statusChip(label: instance.state.label.uppercased(), color: statusColor) } } if let progress = instance.downloadProgress { GeometryReader { geometry in HStack { Spacer() downloadProgressBar(progress: progress) .frame(width: geometry.size.width * 0.5) } } .frame(height: 4) .padding(.top, -8) .padding(.bottom, 2) HStack(spacing: 8) { Text(instance.sharding ?? "") .font(.caption2) .foregroundColor(.secondary) Spacer() downloadSpeedView(progress: progress) } } else { Text(instance.sharding ?? "") .font(.caption2) .foregroundColor(.secondary) } if !instance.chatTasks.isEmpty { VStack(alignment: .leading, spacing: 4) { HStack { Text("Chat Tasks") .font(.caption2) .foregroundColor(.secondary) Text("(\(instance.chatTasks.count))") .font(.caption2) .foregroundColor(.secondary) Spacer() collapseButton(isExpanded: $showChatTasks) } .animation(nil, value: showChatTasks) if showChatTasks { VStack(alignment: .leading, spacing: 8) { ForEach(instance.chatTasks) { task in taskRow(for: task, parentModelName: instance.modelName) } } .transition(.opacity) } } .padding(.top, 4) .animation(.easeInOut(duration: 0.25), value: showChatTasks) } } .padding(.vertical, 6) } private var statusColor: Color { switch instance.state { case .downloading: return .blue case .warmingUp: return .orange case .running: return .green case .ready: return .teal case .waiting, .idle: return .gray case .failed: return .red case .preparing: return .secondary } } @ViewBuilder private func taskRow(for task: InstanceTaskViewModel, parentModelName: String) -> some View { VStack(alignment: .leading, spacing: 4) { HStack(alignment: .top, spacing: 8) { taskStatusIcon(for: task) VStack(alignment: .leading, spacing: 2) { Text("Chat") .font(.caption) .fontWeight(.semibold) if let subtitle = task.subtitle, subtitle.caseInsensitiveCompare(parentModelName) != .orderedSame { Text(subtitle) .font(.caption2) .foregroundColor(.secondary) } if let prompt = task.promptPreview, !prompt.isEmpty { Text("⊙ \(prompt)") .font(.caption2) .foregroundColor(.secondary) .lineLimit(2) } if task.status == .failed, let error = task.errorMessage, !error.isEmpty { Text(error) .font(.caption2) .foregroundColor(.red) .lineLimit(3) } } Spacer(minLength: 6) Button { infoTask = task } label: { Image(systemName: "info.circle") .imageScale(.small) } .buttonStyle(.plain) .popover( item: Binding( get: { infoTask?.id == task.id ? infoTask : nil }, set: { newValue in if newValue == nil { infoTask = nil } else { infoTask = newValue } } ), attachmentAnchor: .rect(.bounds), arrowEdge: .top ) { _ in TaskDetailView(task: task) .padding() .frame(width: 240) } } } } private func taskStatusIcon(for task: InstanceTaskViewModel) -> some View { let icon: String let color: Color let animation: Animation? switch task.status { case .running: icon = "arrow.triangle.2.circlepath" color = .blue animation = Animation.linear(duration: 1).repeatForever(autoreverses: false) case .pending: icon = "circle.dashed" color = .secondary animation = nil case .failed: icon = "exclamationmark.triangle.fill" color = .red animation = nil case .complete: icon = "checkmark.circle.fill" color = .green animation = nil case .unknown: icon = "questionmark.circle" color = .secondary animation = nil } let image = Image(systemName: icon) .imageScale(.small) .foregroundColor(color) if let animation { return AnyView( image .rotationEffect(.degrees(animatedTaskIDs.contains(task.id) ? 360 : 0)) .onAppear { if !animatedTaskIDs.contains(task.id) { animatedTaskIDs.insert(task.id) } } .animation(animation, value: animatedTaskIDs) ) } return AnyView(image) } private func statusChip(label: String, color: Color) -> some View { Text(label) .font(.caption2) .padding(.horizontal, 8) .padding(.vertical, 4) .background(color.opacity(0.15)) .foregroundColor(color) .clipShape(Capsule()) } private func downloadStatusView(progress: DownloadProgressViewModel) -> some View { VStack(alignment: .trailing, spacing: 4) { statusChip(label: "DOWNLOADING", color: .blue) Text(progress.formattedProgress) .foregroundColor(.primary) } .font(.caption2) } private func downloadSpeedView(progress: DownloadProgressViewModel) -> some View { HStack(spacing: 4) { Text(progress.formattedSpeed) if let eta = progress.formattedETA { Text("·") Text(eta) } } .font(.caption2) .foregroundColor(.secondary) } private func downloadProgressBar(progress: DownloadProgressViewModel) -> some View { ProgressView(value: progress.fractionCompleted) .progressViewStyle(.linear) .tint(.blue) } private func collapseButton(isExpanded: Binding) -> some View { Button { isExpanded.wrappedValue.toggle() } label: { Label( isExpanded.wrappedValue ? "Hide" : "Show", systemImage: isExpanded.wrappedValue ? "chevron.up" : "chevron.down" ) .labelStyle(.titleAndIcon) .contentTransition(.symbolEffect(.replace)) } .buttonStyle(.plain) .font(.caption2) } private struct TaskDetailView: View, Identifiable { let task: InstanceTaskViewModel var id: String { task.id } var body: some View { ScrollView { VStack(alignment: .leading, spacing: 12) { parameterSection messageSection if let error = task.errorMessage, !error.isEmpty { detailRow( icon: "exclamationmark.triangle.fill", title: "Error", value: error, tint: .red ) } } } } @ViewBuilder private var parameterSection: some View { if let params = task.parameters { VStack(alignment: .leading, spacing: 6) { Text("Parameters") .font(.subheadline) if let temperature = params.temperature { detailRow(title: "Temperature", value: String(format: "%.1f", temperature)) } if let maxTokens = params.maxTokens { detailRow(title: "Max Tokens", value: "\(maxTokens)") } if let stream = params.stream { detailRow(title: "Stream", value: stream ? "On" : "Off") } if let topP = params.topP { detailRow(title: "Top P", value: String(format: "%.2f", topP)) } } } } @ViewBuilder private var messageSection: some View { if let messages = task.parameters?.messages, !messages.isEmpty { VStack(alignment: .leading, spacing: 6) { Text("Messages") .font(.subheadline) ForEach(Array(messages.enumerated()), id: \.offset) { _, message in VStack(alignment: .leading, spacing: 2) { Text(message.role?.capitalized ?? "Message") .font(.caption) .foregroundColor(.secondary) if let content = message.content, !content.isEmpty { Text(content) .font(.caption2) .foregroundColor(.primary) } } .padding(8) .background(Color.secondary.opacity(0.08)) .clipShape(RoundedRectangle(cornerRadius: 6)) } } } } @ViewBuilder private func detailRow( icon: String? = nil, title: String, value: String, tint: Color = .secondary ) -> some View { HStack(alignment: .firstTextBaseline, spacing: 6) { if let icon { Image(systemName: icon) .imageScale(.small) .foregroundColor(tint) } Text(title) .font(.caption) .foregroundColor(.secondary) Spacer() Text(value) .font(.caption2) .foregroundColor(.primary) } } } } ================================================ FILE: app/EXO/EXO/Views/NodeDetailView.swift ================================================ import SwiftUI struct NodeDetailView: View { let node: NodeViewModel var body: some View { VStack(alignment: .leading, spacing: 10) { Text(node.friendlyName) .font(.headline) Text(node.model) .font(.caption) .foregroundColor(.secondary) Divider() metricRow(label: "Memory", value: node.memoryLabel) ProgressView(value: node.memoryProgress) metricRow(label: "CPU Usage", value: node.cpuUsageLabel) metricRow(label: "GPU Usage", value: node.gpuUsageLabel) metricRow(label: "Temperature", value: node.temperatureLabel) metricRow(label: "Power", value: node.powerLabel) } .padding() } private func metricRow(label: String, value: String) -> some View { HStack { Text(label) .font(.caption) .foregroundColor(.secondary) Spacer() Text(value) .font(.subheadline) } } } ================================================ FILE: app/EXO/EXO/Views/NodeRowView.swift ================================================ import SwiftUI struct NodeRowView: View { let node: NodeViewModel var body: some View { VStack(alignment: .leading, spacing: 4) { HStack { VStack(alignment: .leading) { Text(node.friendlyName) .font(.subheadline) Text(node.memoryLabel) .font(.caption) .foregroundColor(.secondary) } Spacer() VStack(alignment: .trailing) { Text("\(node.gpuUsagePercent, specifier: "%.0f")% GPU") .font(.caption) Text(node.temperatureLabel) .font(.caption2) .foregroundColor(.secondary) } } ProgressView(value: node.memoryProgress) .progressViewStyle(.linear) } .padding(.vertical, 4) } } ================================================ FILE: app/EXO/EXO/Views/SettingsView.swift ================================================ import AppKit import SwiftUI /// Native macOS Settings window following Apple HIG. /// Organized into General, Model, Advanced, and About sections. struct SettingsView: View { @EnvironmentObject private var controller: ExoProcessController @EnvironmentObject private var updater: SparkleUpdater @EnvironmentObject private var networkStatusService: NetworkStatusService @EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService @EnvironmentObject private var stateService: ClusterStateService @State private var pendingNamespace: String = "" @State private var pendingHFToken: String = "" @State private var pendingEnableImageModels = false @State private var pendingOfflineMode = false @State private var needsRestart = false @State private var bugReportInFlight = false @State private var bugReportMessage: String? @State private var uninstallInProgress = false var body: some View { TabView { generalTab .tabItem { Label("General", systemImage: "gear") } modelTab .tabItem { Label("Model", systemImage: "cube") } advancedTab .tabItem { Label("Advanced", systemImage: "wrench.and.screwdriver") } aboutTab .tabItem { Label("About", systemImage: "info.circle") } } .frame(width: 450, height: 400) .onAppear { pendingNamespace = controller.customNamespace pendingHFToken = controller.hfToken pendingEnableImageModels = controller.enableImageModels pendingOfflineMode = controller.offlineMode needsRestart = false } } // MARK: - General Tab private var generalTab: some View { Form { Section { LabeledContent("Cluster Namespace") { TextField("default", text: $pendingNamespace) .textFieldStyle(.roundedBorder) .frame(width: 200) } Text("Nodes with the same namespace form a cluster. Leave empty for default.") .font(.caption) .foregroundColor(.secondary) } Section { LabeledContent("HuggingFace Token") { SecureField("optional", text: $pendingHFToken) .textFieldStyle(.roundedBorder) .frame(width: 200) } Text("Required for gated models. Get yours at huggingface.co/settings/tokens") .font(.caption) .foregroundColor(.secondary) } Section { Toggle("Offline Mode", isOn: $pendingOfflineMode) Text("Skip internet checks and use only locally available models.") .font(.caption) .foregroundColor(.secondary) } Section { HStack { Spacer() Button("Save & Restart") { applyGeneralSettings() } .disabled(!hasGeneralChanges) } } } .formStyle(.grouped) .padding() } // MARK: - Model Tab private var modelTab: some View { Form { Section { Toggle("Enable Image Models (experimental)", isOn: $pendingEnableImageModels) Text("Allow text-to-image and image-to-image models in the model picker.") .font(.caption) .foregroundColor(.secondary) } Section { HStack { Spacer() Button("Save & Restart") { applyModelSettings() } .disabled(!hasModelChanges) } } } .formStyle(.grouped) .padding() } // MARK: - Advanced Tab private var advancedTab: some View { Form { Section("Onboarding") { HStack { VStack(alignment: .leading) { Text("Reset Onboarding") Text("Opens the dashboard and resets the onboarding wizard.") .font(.caption) .foregroundColor(.secondary) } Spacer() Button("Reset") { guard let url = URL(string: "http://localhost:52415/?reset-onboarding") else { return } NSWorkspace.shared.open(url) } } } Section("Debug Info") { LabeledContent("Thunderbolt Bridge") { Text(thunderboltStatusText) .foregroundColor(thunderboltStatusColor) } VStack(alignment: .leading, spacing: 2) { clusterThunderboltBridgeView } VStack(alignment: .leading, spacing: 2) { interfaceIpList } VStack(alignment: .leading, spacing: 2) { rdmaStatusView } sendBugReportButton } Section("Danger Zone") { Button(role: .destructive) { showUninstallConfirmationAlert() } label: { HStack { Text("Uninstall EXO") Spacer() Image(systemName: "trash") .imageScale(.small) } } .disabled(uninstallInProgress) } } .formStyle(.grouped) .padding() } // MARK: - About Tab private var aboutTab: some View { Form { Section { LabeledContent("Version") { Text(buildTag) .textSelection(.enabled) } LabeledContent("Commit") { Text(buildCommit) .font(.system(.body, design: .monospaced)) .textSelection(.enabled) } } Section { Button("Check for Updates") { updater.checkForUpdates() } } } .formStyle(.grouped) .padding() } // MARK: - Debug Info Views (moved from ContentView) private var thunderboltStatusText: String { switch networkStatusService.status.thunderboltBridgeState { case .some(.disabled): return "Disabled" case .some(.deleted): return "Deleted" case .some(.enabled): return "Enabled" case nil: return "Unknown" } } private var thunderboltStatusColor: Color { switch networkStatusService.status.thunderboltBridgeState { case .some(.disabled), .some(.deleted): return .green case .some(.enabled): return .red case nil: return .secondary } } private var clusterThunderboltBridgeView: some View { let bridgeStatuses = stateService.latestSnapshot?.nodeThunderboltBridge ?? [:] let localNodeId = stateService.localNodeId let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:] return VStack(alignment: .leading, spacing: 1) { if bridgeStatuses.isEmpty { Text("Cluster TB Bridge: No data") .font(.caption2) .foregroundColor(.secondary) } else { Text("Cluster TB Bridge Status:") .font(.caption2) .foregroundColor(.secondary) ForEach(Array(bridgeStatuses.keys.sorted()), id: \.self) { nodeId in if let status = bridgeStatuses[nodeId] { let nodeName = nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8)) let isLocal = nodeId == localNodeId let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):" let statusText = !status.exists ? "N/A" : (status.enabled ? "Enabled" : "Disabled") let color: Color = !status.exists ? .secondary : (status.enabled ? .red : .green) Text("\(prefix) \(statusText)") .font(.caption2) .foregroundColor(color) } } } } } private var interfaceIpList: some View { let statuses = networkStatusService.status.interfaceStatuses return VStack(alignment: .leading, spacing: 1) { Text("Interfaces (en0–en7):") .font(.caption2) .foregroundColor(.secondary) if statuses.isEmpty { Text(" Unknown") .font(.caption2) .foregroundColor(.secondary) } else { ForEach(statuses, id: \.interfaceName) { status in let ipText = status.ipAddress ?? "No IP" Text(" \(status.interfaceName): \(ipText)") .font(.caption2) .foregroundColor(status.ipAddress == nil ? .red : .green) } } } } private var rdmaStatusView: some View { let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:] let localNodeId = stateService.localNodeId let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:] let localDevices = networkStatusService.status.localRdmaDevices let localPorts = networkStatusService.status.localRdmaActivePorts return VStack(alignment: .leading, spacing: 1) { if rdmaStatuses.isEmpty { Text("Cluster RDMA: No data") .font(.caption2) .foregroundColor(.secondary) } else { Text("Cluster RDMA Status:") .font(.caption2) .foregroundColor(.secondary) ForEach(Array(rdmaStatuses.keys.sorted()), id: \.self) { nodeId in if let status = rdmaStatuses[nodeId] { let nodeName = nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8)) let isLocal = nodeId == localNodeId let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):" let statusText = status.enabled ? "Enabled" : "Disabled" let color: Color = status.enabled ? .green : .orange Text("\(prefix) \(statusText)") .font(.caption2) .foregroundColor(color) } } } if !localDevices.isEmpty { Text(" Local Devices: \(localDevices.joined(separator: ", "))") .font(.caption2) .foregroundColor(.secondary) } if !localPorts.isEmpty { Text(" Local Active Ports:") .font(.caption2) .foregroundColor(.secondary) ForEach(localPorts, id: \.device) { port in Text(" \(port.device) port \(port.port): \(port.state)") .font(.caption2) .foregroundColor(.green) } } } } private var sendBugReportButton: some View { VStack(alignment: .leading, spacing: 4) { Button { Task { await sendBugReport() } } label: { HStack { if bugReportInFlight { ProgressView() .scaleEffect(0.6) } Text("Send Bug Report") .font(.caption) .fontWeight(.semibold) Spacer() } } .disabled(bugReportInFlight) if let message = bugReportMessage { Text(message) .font(.caption2) .foregroundColor(.secondary) .fixedSize(horizontal: false, vertical: true) } } } // MARK: - Actions private func sendBugReport() async { bugReportInFlight = true bugReportMessage = "Collecting logs..." let service = BugReportService() do { let outcome = try await service.sendReport(isManual: true) bugReportMessage = outcome.message } catch { bugReportMessage = error.localizedDescription } bugReportInFlight = false } private func showUninstallConfirmationAlert() { let alert = NSAlert() alert.messageText = "Uninstall EXO" alert.informativeText = """ This will remove EXO and all its system components: • Network configuration daemon • Launch at login registration • EXO network location The app will be moved to Trash. """ alert.alertStyle = .warning alert.addButton(withTitle: "Uninstall") alert.addButton(withTitle: "Cancel") if let uninstallButton = alert.buttons.first { uninstallButton.hasDestructiveAction = true } let response = alert.runModal() if response == .alertFirstButtonReturn { performUninstall() } } private func performUninstall() { uninstallInProgress = true controller.cancelPendingLaunch() controller.stop() stateService.stopPolling() DispatchQueue.global(qos: .utility).async { do { try NetworkSetupHelper.uninstall() DispatchQueue.main.async { LaunchAtLoginHelper.disable() self.moveAppToTrash() DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) { NSApplication.shared.terminate(nil) } } } catch { DispatchQueue.main.async { let errorAlert = NSAlert() errorAlert.messageText = "Uninstall Failed" errorAlert.informativeText = error.localizedDescription errorAlert.alertStyle = .critical errorAlert.addButton(withTitle: "OK") errorAlert.runModal() self.uninstallInProgress = false } } } } private func moveAppToTrash() { guard let appURL = Bundle.main.bundleURL as URL? else { return } do { try FileManager.default.trashItem(at: appURL, resultingItemURL: nil) } catch { // If we can't trash the app, that's OK - user can do it manually } } // MARK: - Helpers private var hasGeneralChanges: Bool { pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken || pendingOfflineMode != controller.offlineMode } private var hasModelChanges: Bool { pendingEnableImageModels != controller.enableImageModels } private func applyGeneralSettings() { controller.customNamespace = pendingNamespace controller.hfToken = pendingHFToken controller.offlineMode = pendingOfflineMode restartIfRunning() } private func applyModelSettings() { controller.enableImageModels = pendingEnableImageModels restartIfRunning() } private func restartIfRunning() { if controller.status == .running || controller.status == .starting { controller.restart() } } private var buildTag: String { Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown" } private var buildCommit: String { Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown" } } ================================================ FILE: app/EXO/EXO/Views/SettingsWindowController.swift ================================================ import AppKit import SwiftUI /// Manages a standalone native macOS Settings window. /// Ensures only one instance exists and brings it to front on repeated opens. @MainActor final class SettingsWindowController: ObservableObject { private var window: NSWindow? func open( controller: ExoProcessController, updater: SparkleUpdater, networkStatusService: NetworkStatusService, thunderboltBridgeService: ThunderboltBridgeService, stateService: ClusterStateService ) { if let existing = window, existing.isVisible { existing.makeKeyAndOrderFront(nil) NSApp.activate(ignoringOtherApps: true) return } let settingsView = SettingsView() .environmentObject(controller) .environmentObject(updater) .environmentObject(networkStatusService) .environmentObject(thunderboltBridgeService) .environmentObject(stateService) let hostingView = NSHostingView(rootView: settingsView) let newWindow = NSWindow( contentRect: NSRect(x: 0, y: 0, width: 450, height: 400), styleMask: [.titled, .closable], backing: .buffered, defer: false ) newWindow.title = "EXO Settings" newWindow.contentView = hostingView newWindow.center() newWindow.isReleasedWhenClosed = false newWindow.makeKeyAndOrderFront(nil) NSApp.activate(ignoringOtherApps: true) window = newWindow } } ================================================ FILE: app/EXO/EXO/Views/TopologyMiniView.swift ================================================ import SwiftUI struct TopologyMiniView: View { let topology: TopologyViewModel var body: some View { VStack(alignment: .leading, spacing: 8) { Text("Topology") .font(.caption) .foregroundColor(.secondary) GeometryReader { geo in ZStack { connectionLines(in: geo.size) let positions = positionedNodes(in: geo.size) ForEach(Array(positions.enumerated()), id: \.element.node.id) { _, positioned in NodeGlyphView(node: positioned.node, isCurrent: positioned.isCurrent) .position(positioned.point) } } } .frame(height: heightForNodes()) } } private func positionedNodes(in size: CGSize) -> [PositionedNode] { let nodes = orderedNodesForLayout() guard !nodes.isEmpty else { return [] } var result: [PositionedNode] = [] let glyphHeight: CGFloat = 70 let rootPoint = CGPoint(x: size.width / 2, y: glyphHeight / 2 + 10) result.append( PositionedNode( node: nodes[0], point: rootPoint, isCurrent: nodes[0].id == topology.currentNodeId ) ) guard nodes.count > 1 else { return result } let childCount = nodes.count - 1 // Larger radius to reduce overlap when several nodes exist let minDimension = min(size.width, size.height) let radius = max(120, minDimension * 0.42) let startAngle = Double.pi * 0.75 let endAngle = Double.pi * 0.25 let step = childCount == 1 ? 0 : (startAngle - endAngle) / Double(childCount - 1) for (index, node) in nodes.dropFirst().enumerated() { let angle = startAngle - step * Double(index) let x = size.width / 2 + radius * CGFloat(cos(angle)) let y = rootPoint.y + radius * CGFloat(sin(angle)) + glyphHeight / 2 result.append( PositionedNode( node: node, point: CGPoint(x: x, y: y), isCurrent: node.id == topology.currentNodeId ) ) } return result } private func orderedNodesForLayout() -> [NodeViewModel] { guard let currentId = topology.currentNodeId else { return topology.nodes } guard let currentIndex = topology.nodes.firstIndex(where: { $0.id == currentId }) else { return topology.nodes } if currentIndex == 0 { return topology.nodes } var reordered = topology.nodes let current = reordered.remove(at: currentIndex) reordered.insert(current, at: 0) return reordered } private func connectionLines(in size: CGSize) -> some View { let positions = positionedNodes(in: size) let positionById = Dictionary( uniqueKeysWithValues: positions.map { ($0.node.id, $0.point) }) return Canvas { context, _ in guard !topology.edges.isEmpty else { return } let nodeRadius: CGFloat = 32 let arrowLength: CGFloat = 10 let arrowSpread: CGFloat = .pi / 7 for edge in topology.edges { guard let start = positionById[edge.sourceId], let end = positionById[edge.targetId] else { continue } let dx = end.x - start.x let dy = end.y - start.y let distance = max(CGFloat(hypot(dx, dy)), 1) let ux = dx / distance let uy = dy / distance let adjustedStart = CGPoint( x: start.x + ux * nodeRadius, y: start.y + uy * nodeRadius) let adjustedEnd = CGPoint(x: end.x - ux * nodeRadius, y: end.y - uy * nodeRadius) var linePath = Path() linePath.move(to: adjustedStart) linePath.addLine(to: adjustedEnd) context.stroke( linePath, with: .color(.secondary.opacity(0.3)), style: StrokeStyle(lineWidth: 1, dash: [4, 4]) ) let angle = atan2(uy, ux) let tip = adjustedEnd let leftWing = CGPoint( x: tip.x - arrowLength * cos(angle - arrowSpread), y: tip.y - arrowLength * sin(angle - arrowSpread) ) let rightWing = CGPoint( x: tip.x - arrowLength * cos(angle + arrowSpread), y: tip.y - arrowLength * sin(angle + arrowSpread) ) var arrowPath = Path() arrowPath.move(to: tip) arrowPath.addLine(to: leftWing) arrowPath.move(to: tip) arrowPath.addLine(to: rightWing) context.stroke( arrowPath, with: .color(.secondary.opacity(0.5)), style: StrokeStyle(lineWidth: 1) ) } } } private func heightForNodes() -> CGFloat { switch topology.nodes.count { case 0...1: return 130 case 2...3: return 200 default: return 240 } } private struct PositionedNode { let node: NodeViewModel let point: CGPoint let isCurrent: Bool } } private struct NodeGlyphView: View { let node: NodeViewModel let isCurrent: Bool var body: some View { VStack(spacing: 2) { Image(systemName: node.deviceIconName) .font(.subheadline) Text(node.friendlyName) .font(.caption2) .lineLimit(1) .foregroundColor(isCurrent ? Color(nsColor: .systemBlue) : .primary) Text(node.memoryLabel) .font(.caption2) HStack(spacing: 3) { Text(node.gpuUsageLabel) Text(node.temperatureLabel) } .foregroundColor(.secondary) .font(.caption2) } .padding(.vertical, 3) .frame(width: 95) } } ================================================ FILE: app/EXO/EXO/main.swift ================================================ // // main.swift // EXO // // Created by Jake Hillion on 2026-02-03. // import Foundation /// Command line options for the EXO app enum CLICommand { case install case uninstall case help case none } /// Parse command line arguments to determine the CLI command func parseArguments() -> CLICommand { let args = CommandLine.arguments if args.contains("--help") || args.contains("-h") { return .help } if args.contains("--install") { return .install } if args.contains("--uninstall") { return .uninstall } return .none } /// Print usage information func printUsage() { let programName = (CommandLine.arguments.first as NSString?)?.lastPathComponent ?? "EXO" print( """ Usage: \(programName) [OPTIONS] Options: --install Install EXO network configuration (requires root) --uninstall Uninstall EXO network configuration (requires root) --help, -h Show this help message When run without options, starts the normal GUI application. Examples: sudo \(programName) --install Install network components as root sudo \(programName) --uninstall Remove network components as root """) } /// Check if running as root func isRunningAsRoot() -> Bool { return getuid() == 0 } // Main entry point let command = parseArguments() switch command { case .help: printUsage() exit(0) case .install: if !isRunningAsRoot() { fputs("Error: --install requires root privileges. Run with sudo.\n", stderr) exit(1) } let success = NetworkSetupHelper.installDirectly() exit(success ? 0 : 1) case .uninstall: if !isRunningAsRoot() { fputs("Error: --uninstall requires root privileges. Run with sudo.\n", stderr) exit(1) } let success = NetworkSetupHelper.uninstallDirectly() exit(success ? 0 : 1) case .none: // Start normal GUI application EXOApp.main() } ================================================ FILE: app/EXO/EXO.xcodeproj/project.pbxproj ================================================ // !$*UTF8*$! { archiveVersion = 1; classes = { }; objectVersion = 77; objects = { /* Begin PBXBuildFile section */ E0140D402ED1F909001F3171 /* exo in Resources */ = {isa = PBXBuildFile; fileRef = E0140D3F2ED1F909001F3171 /* exo */; }; E0A1B1002F5A000100000003 /* Sparkle in Frameworks */ = {isa = PBXBuildFile; productRef = E0A1B1002F5A000100000002 /* Sparkle */; }; /* End PBXBuildFile section */ /* Begin PBXContainerItemProxy section */ E0140D212ED1F79B001F3171 /* PBXContainerItemProxy */ = { isa = PBXContainerItemProxy; containerPortal = E0140D072ED1F79A001F3171 /* Project object */; proxyType = 1; remoteGlobalIDString = E0140D0E2ED1F79A001F3171; remoteInfo = EXO; }; E0140D2B2ED1F79B001F3171 /* PBXContainerItemProxy */ = { isa = PBXContainerItemProxy; containerPortal = E0140D072ED1F79A001F3171 /* Project object */; proxyType = 1; remoteGlobalIDString = E0140D0E2ED1F79A001F3171; remoteInfo = EXO; }; /* End PBXContainerItemProxy section */ /* Begin PBXFileReference section */ E0140D0F2ED1F79A001F3171 /* EXO.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = EXO.app; sourceTree = BUILT_PRODUCTS_DIR; }; E0140D202ED1F79B001F3171 /* EXOTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = EXOTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; }; E0140D2A2ED1F79B001F3171 /* EXOUITests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = EXOUITests.xctest; sourceTree = BUILT_PRODUCTS_DIR; }; E0140D3F2ED1F909001F3171 /* exo */ = {isa = PBXFileReference; lastKnownFileType = folder; name = exo; path = ../../dist/exo; sourceTree = ""; }; /* End PBXFileReference section */ /* Begin PBXFileSystemSynchronizedRootGroup section */ E0140D112ED1F79A001F3171 /* EXO */ = { isa = PBXFileSystemSynchronizedRootGroup; path = EXO; sourceTree = ""; }; E0140D232ED1F79B001F3171 /* EXOTests */ = { isa = PBXFileSystemSynchronizedRootGroup; path = EXOTests; sourceTree = ""; }; E0140D2D2ED1F79B001F3171 /* EXOUITests */ = { isa = PBXFileSystemSynchronizedRootGroup; path = EXOUITests; sourceTree = ""; }; /* End PBXFileSystemSynchronizedRootGroup section */ /* Begin PBXFrameworksBuildPhase section */ E0140D0C2ED1F79A001F3171 /* Frameworks */ = { isa = PBXFrameworksBuildPhase; buildActionMask = 2147483647; files = ( E0A1B1002F5A000100000003 /* Sparkle in Frameworks */, ); runOnlyForDeploymentPostprocessing = 0; }; E0140D1D2ED1F79B001F3171 /* Frameworks */ = { isa = PBXFrameworksBuildPhase; buildActionMask = 2147483647; files = ( ); runOnlyForDeploymentPostprocessing = 0; }; E0140D272ED1F79B001F3171 /* Frameworks */ = { isa = PBXFrameworksBuildPhase; buildActionMask = 2147483647; files = ( ); runOnlyForDeploymentPostprocessing = 0; }; /* End PBXFrameworksBuildPhase section */ /* Begin PBXGroup section */ E0140D062ED1F79A001F3171 = { isa = PBXGroup; children = ( E0140D3F2ED1F909001F3171 /* exo */, E0140D112ED1F79A001F3171 /* EXO */, E0140D232ED1F79B001F3171 /* EXOTests */, E0140D2D2ED1F79B001F3171 /* EXOUITests */, E0140D102ED1F79A001F3171 /* Products */, ); sourceTree = ""; }; E0140D102ED1F79A001F3171 /* Products */ = { isa = PBXGroup; children = ( E0140D0F2ED1F79A001F3171 /* EXO.app */, E0140D202ED1F79B001F3171 /* EXOTests.xctest */, E0140D2A2ED1F79B001F3171 /* EXOUITests.xctest */, ); name = Products; sourceTree = ""; }; /* End PBXGroup section */ /* Begin PBXNativeTarget section */ E0140D0E2ED1F79A001F3171 /* EXO */ = { isa = PBXNativeTarget; buildConfigurationList = E0140D342ED1F79B001F3171 /* Build configuration list for PBXNativeTarget "EXO" */; buildPhases = ( E0140D0B2ED1F79A001F3171 /* Sources */, E0140D0C2ED1F79A001F3171 /* Frameworks */, E0140D0D2ED1F79A001F3171 /* Resources */, ); buildRules = ( ); dependencies = ( ); fileSystemSynchronizedGroups = ( E0140D112ED1F79A001F3171 /* EXO */, ); name = EXO; packageProductDependencies = ( E0A1B1002F5A000100000002 /* Sparkle */, ); productName = EXO; productReference = E0140D0F2ED1F79A001F3171 /* EXO.app */; productType = "com.apple.product-type.application"; }; E0140D1F2ED1F79B001F3171 /* EXOTests */ = { isa = PBXNativeTarget; buildConfigurationList = E0140D372ED1F79B001F3171 /* Build configuration list for PBXNativeTarget "EXOTests" */; buildPhases = ( E0140D1C2ED1F79B001F3171 /* Sources */, E0140D1D2ED1F79B001F3171 /* Frameworks */, E0140D1E2ED1F79B001F3171 /* Resources */, ); buildRules = ( ); dependencies = ( E0140D222ED1F79B001F3171 /* PBXTargetDependency */, ); fileSystemSynchronizedGroups = ( E0140D232ED1F79B001F3171 /* EXOTests */, ); name = EXOTests; packageProductDependencies = ( ); productName = EXOTests; productReference = E0140D202ED1F79B001F3171 /* EXOTests.xctest */; productType = "com.apple.product-type.bundle.unit-test"; }; E0140D292ED1F79B001F3171 /* EXOUITests */ = { isa = PBXNativeTarget; buildConfigurationList = E0140D3A2ED1F79B001F3171 /* Build configuration list for PBXNativeTarget "EXOUITests" */; buildPhases = ( E0140D262ED1F79B001F3171 /* Sources */, E0140D272ED1F79B001F3171 /* Frameworks */, E0140D282ED1F79B001F3171 /* Resources */, ); buildRules = ( ); dependencies = ( E0140D2C2ED1F79B001F3171 /* PBXTargetDependency */, ); fileSystemSynchronizedGroups = ( E0140D2D2ED1F79B001F3171 /* EXOUITests */, ); name = EXOUITests; packageProductDependencies = ( ); productName = EXOUITests; productReference = E0140D2A2ED1F79B001F3171 /* EXOUITests.xctest */; productType = "com.apple.product-type.bundle.ui-testing"; }; /* End PBXNativeTarget section */ /* Begin PBXProject section */ E0140D072ED1F79A001F3171 /* Project object */ = { isa = PBXProject; attributes = { BuildIndependentTargetsInParallel = 1; LastSwiftUpdateCheck = 1610; LastUpgradeCheck = 1610; TargetAttributes = { E0140D0E2ED1F79A001F3171 = { CreatedOnToolsVersion = 16.1; }; E0140D1F2ED1F79B001F3171 = { CreatedOnToolsVersion = 16.1; TestTargetID = E0140D0E2ED1F79A001F3171; }; E0140D292ED1F79B001F3171 = { CreatedOnToolsVersion = 16.1; TestTargetID = E0140D0E2ED1F79A001F3171; }; }; }; buildConfigurationList = E0140D0A2ED1F79A001F3171 /* Build configuration list for PBXProject "EXO" */; developmentRegion = en; hasScannedForEncodings = 0; knownRegions = ( en, Base, ); mainGroup = E0140D062ED1F79A001F3171; minimizedProjectReferenceProxies = 1; packageReferences = ( E0A1B1002F5A000100000001 /* XCRemoteSwiftPackageReference "Sparkle" */, ); preferredProjectObjectVersion = 77; productRefGroup = E0140D102ED1F79A001F3171 /* Products */; projectDirPath = ""; projectRoot = ""; targets = ( E0140D0E2ED1F79A001F3171 /* EXO */, E0140D1F2ED1F79B001F3171 /* EXOTests */, E0140D292ED1F79B001F3171 /* EXOUITests */, ); }; /* End PBXProject section */ /* Begin PBXResourcesBuildPhase section */ E0140D0D2ED1F79A001F3171 /* Resources */ = { isa = PBXResourcesBuildPhase; buildActionMask = 2147483647; files = ( E0140D402ED1F909001F3171 /* exo in Resources */, ); runOnlyForDeploymentPostprocessing = 0; }; E0140D1E2ED1F79B001F3171 /* Resources */ = { isa = PBXResourcesBuildPhase; buildActionMask = 2147483647; files = ( ); runOnlyForDeploymentPostprocessing = 0; }; E0140D282ED1F79B001F3171 /* Resources */ = { isa = PBXResourcesBuildPhase; buildActionMask = 2147483647; files = ( ); runOnlyForDeploymentPostprocessing = 0; }; /* End PBXResourcesBuildPhase section */ /* Begin PBXSourcesBuildPhase section */ E0140D0B2ED1F79A001F3171 /* Sources */ = { isa = PBXSourcesBuildPhase; buildActionMask = 2147483647; files = ( ); runOnlyForDeploymentPostprocessing = 0; }; E0140D1C2ED1F79B001F3171 /* Sources */ = { isa = PBXSourcesBuildPhase; buildActionMask = 2147483647; files = ( ); runOnlyForDeploymentPostprocessing = 0; }; E0140D262ED1F79B001F3171 /* Sources */ = { isa = PBXSourcesBuildPhase; buildActionMask = 2147483647; files = ( ); runOnlyForDeploymentPostprocessing = 0; }; /* End PBXSourcesBuildPhase section */ /* Begin PBXTargetDependency section */ E0140D222ED1F79B001F3171 /* PBXTargetDependency */ = { isa = PBXTargetDependency; target = E0140D0E2ED1F79A001F3171 /* EXO */; targetProxy = E0140D212ED1F79B001F3171 /* PBXContainerItemProxy */; }; E0140D2C2ED1F79B001F3171 /* PBXTargetDependency */ = { isa = PBXTargetDependency; target = E0140D0E2ED1F79A001F3171 /* EXO */; targetProxy = E0140D2B2ED1F79B001F3171 /* PBXContainerItemProxy */; }; /* End PBXTargetDependency section */ /* Begin XCBuildConfiguration section */ E0140D322ED1F79B001F3171 /* Debug */ = { isa = XCBuildConfiguration; buildSettings = { ALWAYS_SEARCH_USER_PATHS = NO; ASSETCATALOG_COMPILER_GENERATE_SWIFT_ASSET_SYMBOL_EXTENSIONS = YES; CLANG_ANALYZER_NONNULL = YES; CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; CLANG_CXX_LANGUAGE_STANDARD = "gnu++20"; CLANG_ENABLE_MODULES = YES; CLANG_ENABLE_OBJC_ARC = YES; CLANG_ENABLE_OBJC_WEAK = YES; CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; CLANG_WARN_BOOL_CONVERSION = YES; CLANG_WARN_COMMA = YES; CLANG_WARN_CONSTANT_CONVERSION = YES; CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; CLANG_WARN_DOCUMENTATION_COMMENTS = YES; CLANG_WARN_EMPTY_BODY = YES; CLANG_WARN_ENUM_CONVERSION = YES; CLANG_WARN_INFINITE_RECURSION = YES; CLANG_WARN_INT_CONVERSION = YES; CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES; CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; CLANG_WARN_STRICT_PROTOTYPES = YES; CLANG_WARN_SUSPICIOUS_MOVE = YES; CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; CLANG_WARN_UNREACHABLE_CODE = YES; CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; COPY_PHASE_STRIP = NO; DEBUG_INFORMATION_FORMAT = dwarf; ENABLE_STRICT_OBJC_MSGSEND = YES; ENABLE_TESTABILITY = YES; ENABLE_USER_SCRIPT_SANDBOXING = YES; GCC_C_LANGUAGE_STANDARD = gnu17; GCC_DYNAMIC_NO_PIC = NO; GCC_NO_COMMON_BLOCKS = YES; GCC_OPTIMIZATION_LEVEL = 0; GCC_PREPROCESSOR_DEFINITIONS = ( "DEBUG=1", "$(inherited)", ); GCC_WARN_64_TO_32_BIT_CONVERSION = YES; GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; GCC_WARN_UNDECLARED_SELECTOR = YES; GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; GCC_WARN_UNUSED_FUNCTION = YES; GCC_WARN_UNUSED_VARIABLE = YES; LOCALIZATION_PREFERS_STRING_CATALOGS = YES; MACOSX_DEPLOYMENT_TARGET = 15.1; MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE; MTL_FAST_MATH = YES; ONLY_ACTIVE_ARCH = YES; SDKROOT = macosx; SWIFT_ACTIVE_COMPILATION_CONDITIONS = "DEBUG $(inherited)"; SWIFT_OPTIMIZATION_LEVEL = "-Onone"; SWIFT_TREAT_WARNINGS_AS_ERRORS = YES; GCC_TREAT_WARNINGS_AS_ERRORS = YES; }; name = Debug; }; E0140D332ED1F79B001F3171 /* Release */ = { isa = XCBuildConfiguration; buildSettings = { ALWAYS_SEARCH_USER_PATHS = NO; ASSETCATALOG_COMPILER_GENERATE_SWIFT_ASSET_SYMBOL_EXTENSIONS = YES; CLANG_ANALYZER_NONNULL = YES; CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; CLANG_CXX_LANGUAGE_STANDARD = "gnu++20"; CLANG_ENABLE_MODULES = YES; CLANG_ENABLE_OBJC_ARC = YES; CLANG_ENABLE_OBJC_WEAK = YES; CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; CLANG_WARN_BOOL_CONVERSION = YES; CLANG_WARN_COMMA = YES; CLANG_WARN_CONSTANT_CONVERSION = YES; CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; CLANG_WARN_DOCUMENTATION_COMMENTS = YES; CLANG_WARN_EMPTY_BODY = YES; CLANG_WARN_ENUM_CONVERSION = YES; CLANG_WARN_INFINITE_RECURSION = YES; CLANG_WARN_INT_CONVERSION = YES; CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES; CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; CLANG_WARN_STRICT_PROTOTYPES = YES; CLANG_WARN_SUSPICIOUS_MOVE = YES; CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; CLANG_WARN_UNREACHABLE_CODE = YES; CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; COPY_PHASE_STRIP = NO; DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym"; ENABLE_NS_ASSERTIONS = NO; ENABLE_STRICT_OBJC_MSGSEND = YES; ENABLE_USER_SCRIPT_SANDBOXING = YES; GCC_C_LANGUAGE_STANDARD = gnu17; GCC_NO_COMMON_BLOCKS = YES; GCC_WARN_64_TO_32_BIT_CONVERSION = YES; GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; GCC_WARN_UNDECLARED_SELECTOR = YES; GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; GCC_WARN_UNUSED_FUNCTION = YES; GCC_WARN_UNUSED_VARIABLE = YES; LOCALIZATION_PREFERS_STRING_CATALOGS = YES; MACOSX_DEPLOYMENT_TARGET = 15.1; MTL_ENABLE_DEBUG_INFO = NO; MTL_FAST_MATH = YES; SDKROOT = macosx; SWIFT_COMPILATION_MODE = wholemodule; SWIFT_TREAT_WARNINGS_AS_ERRORS = YES; GCC_TREAT_WARNINGS_AS_ERRORS = YES; }; name = Release; }; E0140D352ED1F79B001F3171 /* Debug */ = { isa = XCBuildConfiguration; buildSettings = { ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor; CODE_SIGN_ENTITLEMENTS = EXO/EXO.entitlements; CODE_SIGN_STYLE = Automatic; COMBINE_HIDPI_IMAGES = YES; CURRENT_PROJECT_VERSION = 1; DEVELOPMENT_ASSET_PATHS = "\"EXO/Preview Content\""; ENABLE_PREVIEWS = YES; GENERATE_INFOPLIST_FILE = YES; INFOPLIST_FILE = EXO/Info.plist; INFOPLIST_KEY_LSUIElement = YES; INFOPLIST_KEY_EXOBuildCommit = "$(EXO_BUILD_COMMIT)"; INFOPLIST_KEY_EXOBuildTag = "$(EXO_BUILD_TAG)"; INFOPLIST_KEY_NSAppleEventsUsageDescription = "EXO needs to run a signed network setup script with administrator privileges."; INFOPLIST_KEY_NSHumanReadableCopyright = ""; INFOPLIST_KEY_SUEnableAutomaticChecks = YES; INFOPLIST_KEY_SUFeedURL = "$(SPARKLE_FEED_URL)"; INFOPLIST_KEY_SUPublicEDKey = "$(SPARKLE_ED25519_PUBLIC)"; LD_RUNPATH_SEARCH_PATHS = ( "$(inherited)", "@executable_path/../Frameworks", ); MARKETING_VERSION = 1.0.1; PRODUCT_BUNDLE_IDENTIFIER = exolabs.EXO; PRODUCT_NAME = "$(TARGET_NAME)"; EXO_BUILD_COMMIT = local; EXO_BUILD_TAG = dev; SPARKLE_ED25519_PUBLIC = ""; SPARKLE_FEED_URL = ""; SWIFT_EMIT_LOC_STRINGS = YES; SWIFT_VERSION = 5.0; }; name = Debug; }; E0140D362ED1F79B001F3171 /* Release */ = { isa = XCBuildConfiguration; buildSettings = { ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor; CODE_SIGN_ENTITLEMENTS = EXO/EXO.entitlements; CODE_SIGN_STYLE = Automatic; COMBINE_HIDPI_IMAGES = YES; CURRENT_PROJECT_VERSION = 1; DEVELOPMENT_ASSET_PATHS = "\"EXO/Preview Content\""; ENABLE_PREVIEWS = YES; GENERATE_INFOPLIST_FILE = YES; INFOPLIST_FILE = EXO/Info.plist; INFOPLIST_KEY_LSUIElement = YES; INFOPLIST_KEY_EXOBuildCommit = "$(EXO_BUILD_COMMIT)"; INFOPLIST_KEY_EXOBuildTag = "$(EXO_BUILD_TAG)"; INFOPLIST_KEY_NSAppleEventsUsageDescription = "EXO needs to run a signed network setup script with administrator privileges."; INFOPLIST_KEY_NSHumanReadableCopyright = ""; INFOPLIST_KEY_SUEnableAutomaticChecks = YES; INFOPLIST_KEY_SUFeedURL = "$(SPARKLE_FEED_URL)"; INFOPLIST_KEY_SUPublicEDKey = "$(SPARKLE_ED25519_PUBLIC)"; LD_RUNPATH_SEARCH_PATHS = ( "$(inherited)", "@executable_path/../Frameworks", ); MARKETING_VERSION = 1.0.1; PRODUCT_BUNDLE_IDENTIFIER = exolabs.EXO; PRODUCT_NAME = "$(TARGET_NAME)"; EXO_BUILD_COMMIT = local; EXO_BUILD_TAG = dev; SPARKLE_ED25519_PUBLIC = ""; SPARKLE_FEED_URL = ""; SWIFT_EMIT_LOC_STRINGS = YES; SWIFT_VERSION = 5.0; }; name = Release; }; E0140D382ED1F79B001F3171 /* Debug */ = { isa = XCBuildConfiguration; buildSettings = { BUNDLE_LOADER = "$(TEST_HOST)"; CODE_SIGN_STYLE = Automatic; CURRENT_PROJECT_VERSION = 1; GENERATE_INFOPLIST_FILE = YES; MACOSX_DEPLOYMENT_TARGET = 15.1; MARKETING_VERSION = 1.0; PRODUCT_BUNDLE_IDENTIFIER = exolabs.EXOTests; PRODUCT_NAME = "$(TARGET_NAME)"; SWIFT_EMIT_LOC_STRINGS = NO; SWIFT_VERSION = 5.0; TEST_HOST = "$(BUILT_PRODUCTS_DIR)/EXO.app/$(BUNDLE_EXECUTABLE_FOLDER_PATH)/EXO"; }; name = Debug; }; E0140D392ED1F79B001F3171 /* Release */ = { isa = XCBuildConfiguration; buildSettings = { BUNDLE_LOADER = "$(TEST_HOST)"; CODE_SIGN_STYLE = Automatic; CURRENT_PROJECT_VERSION = 1; GENERATE_INFOPLIST_FILE = YES; MACOSX_DEPLOYMENT_TARGET = 15.1; MARKETING_VERSION = 1.0; PRODUCT_BUNDLE_IDENTIFIER = exolabs.EXOTests; PRODUCT_NAME = "$(TARGET_NAME)"; SWIFT_EMIT_LOC_STRINGS = NO; SWIFT_VERSION = 5.0; TEST_HOST = "$(BUILT_PRODUCTS_DIR)/EXO.app/$(BUNDLE_EXECUTABLE_FOLDER_PATH)/EXO"; }; name = Release; }; E0140D3B2ED1F79B001F3171 /* Debug */ = { isa = XCBuildConfiguration; buildSettings = { CODE_SIGN_STYLE = Automatic; CURRENT_PROJECT_VERSION = 1; GENERATE_INFOPLIST_FILE = YES; MARKETING_VERSION = 1.0; PRODUCT_BUNDLE_IDENTIFIER = exolabs.EXOUITests; PRODUCT_NAME = "$(TARGET_NAME)"; SWIFT_EMIT_LOC_STRINGS = NO; SWIFT_VERSION = 5.0; TEST_TARGET_NAME = EXO; }; name = Debug; }; E0140D3C2ED1F79B001F3171 /* Release */ = { isa = XCBuildConfiguration; buildSettings = { CODE_SIGN_STYLE = Automatic; CURRENT_PROJECT_VERSION = 1; GENERATE_INFOPLIST_FILE = YES; MARKETING_VERSION = 1.0; PRODUCT_BUNDLE_IDENTIFIER = exolabs.EXOUITests; PRODUCT_NAME = "$(TARGET_NAME)"; SWIFT_EMIT_LOC_STRINGS = NO; SWIFT_VERSION = 5.0; TEST_TARGET_NAME = EXO; }; name = Release; }; /* End XCBuildConfiguration section */ /* Begin XCConfigurationList section */ E0140D0A2ED1F79A001F3171 /* Build configuration list for PBXProject "EXO" */ = { isa = XCConfigurationList; buildConfigurations = ( E0140D322ED1F79B001F3171 /* Debug */, E0140D332ED1F79B001F3171 /* Release */, ); defaultConfigurationIsVisible = 0; defaultConfigurationName = Release; }; E0140D342ED1F79B001F3171 /* Build configuration list for PBXNativeTarget "EXO" */ = { isa = XCConfigurationList; buildConfigurations = ( E0140D352ED1F79B001F3171 /* Debug */, E0140D362ED1F79B001F3171 /* Release */, ); defaultConfigurationIsVisible = 0; defaultConfigurationName = Release; }; E0140D372ED1F79B001F3171 /* Build configuration list for PBXNativeTarget "EXOTests" */ = { isa = XCConfigurationList; buildConfigurations = ( E0140D382ED1F79B001F3171 /* Debug */, E0140D392ED1F79B001F3171 /* Release */, ); defaultConfigurationIsVisible = 0; defaultConfigurationName = Release; }; E0140D3A2ED1F79B001F3171 /* Build configuration list for PBXNativeTarget "EXOUITests" */ = { isa = XCConfigurationList; buildConfigurations = ( E0140D3B2ED1F79B001F3171 /* Debug */, E0140D3C2ED1F79B001F3171 /* Release */, ); defaultConfigurationIsVisible = 0; defaultConfigurationName = Release; }; /* End XCConfigurationList section */ /* Begin XCRemoteSwiftPackageReference section */ E0A1B1002F5A000100000001 /* XCRemoteSwiftPackageReference "Sparkle" */ = { isa = XCRemoteSwiftPackageReference; repositoryURL = "https://github.com/sparkle-project/Sparkle.git"; requirement = { kind = upToNextMajorVersion; minimumVersion = 2.9.0-beta.1; }; }; /* End XCRemoteSwiftPackageReference section */ /* Begin XCSwiftPackageProductDependency section */ E0A1B1002F5A000100000002 /* Sparkle */ = { isa = XCSwiftPackageProductDependency; package = E0A1B1002F5A000100000001 /* XCRemoteSwiftPackageReference "Sparkle" */; productName = Sparkle; }; /* End XCSwiftPackageProductDependency section */ }; rootObject = E0140D072ED1F79A001F3171 /* Project object */; } ================================================ FILE: app/EXO/EXO.xcodeproj/project.xcworkspace/contents.xcworkspacedata ================================================ ================================================ FILE: app/EXO/EXO.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved ================================================ { "originHash" : "5751fcbe53b64441ed73aceb16987d6b3fc3ebc666cb9ec2de1f6a2d441f2515", "pins" : [ { "identity" : "sparkle", "kind" : "remoteSourceControl", "location" : "https://github.com/sparkle-project/Sparkle.git", "state" : { "revision" : "e641adb41915a8409895e2e30666aa64e487b637", "version" : "2.9.0-beta.1" } } ], "version" : 3 } ================================================ FILE: app/EXO/EXO.xcodeproj/xcshareddata/xcschemes/EXO.xcscheme ================================================ ================================================ FILE: app/EXO/EXOTests/EXOTests.swift ================================================ // // EXOTests.swift // EXOTests // // Created by Sami Khan on 2025-11-22. // import Testing @testable import EXO struct EXOTests { @Test func example() async throws { // Write your test here and use APIs like `#expect(...)` to check expected conditions. } } ================================================ FILE: app/EXO/EXOUITests/EXOUITests.swift ================================================ // // EXOUITests.swift // EXOUITests // // Created by Sami Khan on 2025-11-22. // import XCTest final class EXOUITests: XCTestCase { override func setUpWithError() throws { // Put setup code here. This method is called before the invocation of each test method in the class. // In UI tests it is usually best to stop immediately when a failure occurs. continueAfterFailure = false // In UI tests it’s important to set the initial state - such as interface orientation - required for your tests before they run. The setUp method is a good place to do this. } override func tearDownWithError() throws { // Put teardown code here. This method is called after the invocation of each test method in the class. } @MainActor func testExample() throws { // UI tests must launch the application that they test. let app = XCUIApplication() app.launch() // Use XCTAssert and related functions to verify your tests produce the correct results. } @MainActor func testLaunchPerformance() throws { if #available(macOS 10.15, iOS 13.0, tvOS 13.0, watchOS 7.0, *) { // This measures how long it takes to launch your application. measure(metrics: [XCTApplicationLaunchMetric()]) { XCUIApplication().launch() } } } } ================================================ FILE: app/EXO/EXOUITests/EXOUITestsLaunchTests.swift ================================================ // // EXOUITestsLaunchTests.swift // EXOUITests // // Created by Sami Khan on 2025-11-22. // import XCTest final class EXOUITestsLaunchTests: XCTestCase { override class var runsForEachTargetApplicationUIConfiguration: Bool { true } override func setUpWithError() throws { continueAfterFailure = false } @MainActor func testLaunch() throws { let app = XCUIApplication() app.launch() // Insert steps here to perform after app launch but before taking a screenshot, // such as logging into a test account or navigating somewhere in the app let attachment = XCTAttachment(screenshot: app.screenshot()) attachment.name = "Launch Screen" attachment.lifetime = .keepAlways add(attachment) } } ================================================ FILE: app/EXO/uninstall-exo.sh ================================================ #!/usr/bin/env bash # # EXO Uninstaller Script # # This script removes all EXO system components that persist after deleting the app. # Run with: sudo ./uninstall-exo.sh # # Components removed: # - LaunchDaemon: /Library/LaunchDaemons/io.exo.networksetup.plist # - Network script: /Library/Application Support/EXO/ # - Log files: /var/log/io.exo.networksetup.* # - Network location: "exo" # - Launch at login registration # set -euo pipefail LABEL="io.exo.networksetup" SCRIPT_DEST="/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh" PLIST_DEST="/Library/LaunchDaemons/io.exo.networksetup.plist" LOG_OUT="/var/log/${LABEL}.log" LOG_ERR="/var/log/${LABEL}.err.log" APP_BUNDLE_ID="io.exo.EXO" # Colors for output RED='\033[0;31m' GREEN='\033[0;32m' YELLOW='\033[1;33m' NC='\033[0m' # No Color echo_info() { echo -e "${GREEN}[INFO]${NC} $1" } echo_warn() { echo -e "${YELLOW}[WARN]${NC} $1" } echo_error() { echo -e "${RED}[ERROR]${NC} $1" } # Check if running as root if [[ $EUID -ne 0 ]]; then echo_error "This script must be run as root (use sudo)" exit 1 fi echo "" echo "========================================" echo " EXO Uninstaller" echo "========================================" echo "" # Unload the LaunchDaemon if running echo_info "Stopping network setup daemon..." if launchctl list | grep -q "$LABEL"; then launchctl bootout system/"$LABEL" 2>/dev/null || true echo_info "Daemon stopped" else echo_warn "Daemon was not running" fi # Remove LaunchDaemon plist if [[ -f $PLIST_DEST ]]; then rm -f "$PLIST_DEST" echo_info "Removed LaunchDaemon plist" else echo_warn "LaunchDaemon plist not found (already removed?)" fi # Remove the script and parent directory if [[ -f $SCRIPT_DEST ]]; then rm -f "$SCRIPT_DEST" echo_info "Removed network setup script" else echo_warn "Network setup script not found (already removed?)" fi # Remove EXO directory if empty if [[ -d "/Library/Application Support/EXO" ]]; then rmdir "/Library/Application Support/EXO" 2>/dev/null && echo_info "Removed EXO support directory" || echo_warn "EXO support directory not empty, leaving in place" fi # Remove log files if [[ -f $LOG_OUT ]] || [[ -f $LOG_ERR ]]; then rm -f "$LOG_OUT" "$LOG_ERR" echo_info "Removed log files" else echo_warn "Log files not found (already removed?)" fi # Switch back to Automatic network location echo_info "Restoring network configuration..." if networksetup -listlocations | grep -q "^Automatic$"; then networksetup -switchtolocation Automatic 2>/dev/null || true echo_info "Switched to Automatic network location" else echo_warn "Automatic network location not found" fi # Delete the exo network location if it exists if networksetup -listlocations | grep -q "^exo$"; then networksetup -deletelocation exo 2>/dev/null || true echo_info "Deleted 'exo' network location" else echo_warn "'exo' network location not found (already removed?)" fi # Re-enable Thunderbolt Bridge if it exists if networksetup -listnetworkservices 2>/dev/null | grep -q "Thunderbolt Bridge"; then networksetup -setnetworkserviceenabled "Thunderbolt Bridge" on 2>/dev/null || true echo_info "Re-enabled Thunderbolt Bridge" fi # Note about launch at login registration # SMAppService-based login items cannot be removed from a shell script. # They can only be unregistered from within the app itself or manually via System Settings. echo_warn "Launch at login must be removed manually:" echo_warn " System Settings → General → Login Items → Remove EXO" # Check if EXO.app exists in common locations APP_FOUND=false for app_path in "/Applications/EXO.app" "$HOME/Applications/EXO.app"; do if [[ -d $app_path ]]; then if [[ $APP_FOUND == false ]]; then echo "" APP_FOUND=true fi echo_warn "EXO.app found at: $app_path" echo_warn "You may want to move it to Trash manually." fi done echo "" echo "========================================" echo_info "EXO uninstall complete!" echo "========================================" echo "" echo "The following have been removed:" echo " • Network setup LaunchDaemon" echo " • Network configuration script" echo " • Log files" echo " • 'exo' network location" echo "" echo "Your network has been restored to use the 'Automatic' location." echo "Thunderbolt Bridge has been re-enabled (if present)." echo "" echo "Manual step required:" echo " Remove EXO from Login Items in System Settings → General → Login Items" echo "" ================================================ FILE: bench/bench.toml ================================================ # Canary benchmark manifest # # Lists the suite files to include. Each file defines benchmarks # with shared constraints, topology, and default args. include = ["single-m3-ultra.toml"] ================================================ FILE: bench/eval_configs/models.toml ================================================ # Model evaluation configurations for exo_eval. # # Each [[model]] entry uses `patterns` — a list of substrings matched # against the model_id. First matching entry wins. # # Required fields: # name, patterns, reasoning # # Optional per-model overrides (CLI flags take priority over these): # temperature, top_p, max_tokens, reasoning_effort # # Fallback defaults (when no per-model config): # reasoning: temperature=1.0, max_tokens=131072, reasoning_effort="high" # non-reasoning: temperature=0.0, max_tokens=16384 # # All per-model values below are sourced from official model cards, # generation_config.json files, and vendor documentation. # ─── Qwen3.5 (Feb 2026) ───────────────────────────────────────────── # Source: HuggingFace model cards (Qwen/Qwen3.5-*) # 35B-A3B thinking general: temp=1.0, top_p=0.95, top_k=20 # 397B thinking: temp=0.6, top_p=0.95, top_k=20 # Non-thinking: temp=0.7, top_p=0.8, top_k=20 # max_tokens: 32768 general, 81920 for complex math/code [[model]] name = "Qwen3.5 2B" patterns = ["Qwen3.5-2B"] reasoning = true temperature = 0.6 top_p = 0.95 max_tokens = 81920 [[model]] name = "Qwen3.5 9B" patterns = ["Qwen3.5-9B"] reasoning = true temperature = 0.6 top_p = 0.95 max_tokens = 81920 [[model]] name = "Qwen3.5 27B" patterns = ["Qwen3.5-27B"] reasoning = true temperature = 0.6 top_p = 0.95 max_tokens = 81920 [[model]] name = "Qwen3.5 35B A3B" patterns = ["Qwen3.5-35B-A3B"] reasoning = true temperature = 1.0 top_p = 0.95 max_tokens = 81920 [[model]] name = "Qwen3.5 122B A10B" patterns = ["Qwen3.5-122B-A10B"] reasoning = true temperature = 0.6 top_p = 0.95 max_tokens = 81920 [[model]] name = "Qwen3.5 397B A17B" patterns = ["Qwen3.5-397B-A17B"] reasoning = true temperature = 0.6 top_p = 0.95 max_tokens = 81920 # ─── Qwen3 (Apr 2025) ─────────────────────────────────────────────── # Source: HuggingFace model cards (Qwen/Qwen3-*) # Thinking: temp=0.6, top_p=0.95, top_k=20 # Non-thinking: temp=0.7, top_p=0.8, top_k=20 # max_tokens: 32768 general, 38912 for complex math/code [[model]] name = "Qwen3 0.6B" patterns = ["Qwen3-0.6B"] reasoning = true temperature = 0.6 top_p = 0.95 max_tokens = 38912 [[model]] name = "Qwen3 30B A3B" patterns = ["Qwen3-30B-A3B"] reasoning = true temperature = 0.6 top_p = 0.95 max_tokens = 38912 [[model]] name = "Qwen3 235B A22B" patterns = ["Qwen3-235B-A22B"] reasoning = true temperature = 0.6 top_p = 0.95 max_tokens = 38912 [[model]] name = "Qwen3 Next 80B Thinking" patterns = ["Qwen3-Next-80B-A3B-Thinking"] reasoning = true temperature = 0.6 top_p = 0.95 max_tokens = 38912 [[model]] name = "Qwen3 Next 80B Instruct" patterns = ["Qwen3-Next-80B-A3B-Instruct"] reasoning = false temperature = 0.7 top_p = 0.8 max_tokens = 16384 [[model]] name = "Qwen3 Coder 480B" patterns = ["Qwen3-Coder-480B"] reasoning = false temperature = 0.7 top_p = 0.8 max_tokens = 16384 [[model]] name = "Qwen3 Coder Next" patterns = ["Qwen3-Coder-Next"] reasoning = false temperature = 0.7 top_p = 0.8 max_tokens = 16384 # ─── GPT-OSS (OpenAI) ─────────────────────────────────────────────── # Source: OpenAI GitHub README + HuggingFace discussion #21 # temp=1.0, top_p=1.0, NO top_k, NO repetition_penalty # reasoning_effort supported: low/medium/high [[model]] name = "GPT-OSS 20B" patterns = ["gpt-oss-20b"] reasoning = true temperature = 1.0 top_p = 1.0 [[model]] name = "GPT-OSS 120B" patterns = ["gpt-oss-120b"] reasoning = true temperature = 1.0 top_p = 1.0 # ─── DeepSeek ──────────────────────────────────────────────────────── # Source: https://api-docs.deepseek.com/quick_start/parameter_settings # Coding/Math: temp=0.0, General: temp=1.3, Creative: temp=1.5 # NOTE: DeepSeek API applies nonlinear temp mapping. These are API values. # When running model directly: API temp 1.0 = model temp 0.3 # We use temp=0.0 for eval (coding/math focus). [[model]] name = "DeepSeek V3.1" patterns = ["DeepSeek-V3.1"] reasoning = true temperature = 0.0 # ─── GLM (ZhipuAI / THUDM) ────────────────────────────────────────── # Source: HuggingFace model cards + generation_config.json + docs.z.ai # GLM 4.5+: temp=1.0, top_p=0.95 # Reasoning tasks: 131072 max_tokens; coding/SWE tasks: temp=0.7 [[model]] name = "GLM-5" patterns = ["GLM-5"] reasoning = true temperature = 1.0 top_p = 0.95 max_tokens = 131072 [[model]] name = "GLM 4.5 Air" patterns = ["GLM-4.5-Air"] reasoning = true temperature = 1.0 top_p = 0.95 [[model]] name = "GLM 4.7" patterns = ["GLM-4.7-"] reasoning = true temperature = 1.0 top_p = 0.95 max_tokens = 131072 # Note: matches both GLM-4.7 and GLM-4.7-Flash # ─── Kimi (Moonshot AI) ───────────────────────────────────────────── # Source: HuggingFace model cards (moonshotai/Kimi-K2-*) # K2-Instruct: temp=0.6 # K2-Thinking: temp=1.0, max_length=262144 # K2.5: thinking temp=1.0, top_p=0.95; instant temp=0.6, top_p=0.95 [[model]] name = "Kimi K2 Thinking" patterns = ["Kimi-K2-Thinking"] reasoning = true temperature = 1.0 max_tokens = 131072 [[model]] name = "Kimi K2.5" patterns = ["Kimi-K2.5"] reasoning = true temperature = 1.0 top_p = 0.95 max_tokens = 131072 [[model]] name = "Kimi K2 Instruct" patterns = ["Kimi-K2-Instruct"] reasoning = false temperature = 0.6 # ─── MiniMax ───────────────────────────────────────────────────────── # Source: HuggingFace model cards + generation_config.json # All models: temp=1.0, top_p=0.95, top_k=40 [[model]] name = "MiniMax M2.5" patterns = ["MiniMax-M2.5"] reasoning = true temperature = 1.0 top_p = 0.95 [[model]] name = "MiniMax M2.1" patterns = ["MiniMax-M2.1"] reasoning = true temperature = 1.0 top_p = 0.95 # ─── Step (StepFun) ───────────────────────────────────────────────── # Source: HuggingFace model card (stepfun-ai/Step-3.5-Flash) # Reasoning: temp=1.0, top_p=0.95 # General chat: temp=0.6, top_p=0.95 # We use reasoning settings for eval. [[model]] name = "Step 3.5 Flash" patterns = ["Step-3.5-Flash"] reasoning = true temperature = 1.0 top_p = 0.95 # ─── Llama (Meta) ─────────────────────────────────────────────────── # Source: generation_config.json + meta-llama/llama-models generation.py # All variants: temp=0.6, top_p=0.9 [[model]] name = "Llama 3.2 1B" patterns = ["Llama-3.2-1B"] reasoning = false temperature = 0.6 top_p = 0.9 [[model]] name = "Llama 3.2 3B" patterns = ["Llama-3.2-3B"] reasoning = false temperature = 0.6 top_p = 0.9 [[model]] name = "Llama 3.1 8B" patterns = ["Llama-3.1-8B", "Meta-Llama-3.1-8B"] reasoning = false temperature = 0.6 top_p = 0.9 [[model]] name = "Llama 3.1 70B" patterns = ["Llama-3.1-70B", "Meta-Llama-3.1-70B"] reasoning = false temperature = 0.6 top_p = 0.9 [[model]] name = "Llama 3.3 70B" patterns = ["Llama-3.3-70B", "llama-3.3-70b"] reasoning = false temperature = 0.6 top_p = 0.9 ================================================ FILE: bench/eval_tool_calls.py ================================================ # pyright: reportAny=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false from __future__ import annotations import argparse import contextlib import json import os import sys import time import tomllib from dataclasses import dataclass, field from pathlib import Path from typing import Any, Literal import httpx from harness import ( ExoClient, ExoHttpError, add_common_instance_args, instance_id_from_instance, nodes_used_in_instance, resolve_model_short_id, run_planning_phase, settle_and_fetch_placements, wait_for_instance_gone, wait_for_instance_ready, ) SCENARIOS_PATH = Path(__file__).parent / "scenarios.toml" @dataclass class Scenario: name: str description: str messages: list[dict[str, Any]] tools: list[dict[str, Any]] expect_tool_call: bool expected_function: str | None = None required_arg_keys: list[str] | None = None tool_result: str | None = None nested_array_key: str | None = None required_item_keys: list[str] | None = None def load_scenarios(path: Path) -> list[Scenario]: with open(path, "rb") as f: data = tomllib.load(f) tools_data = data.get("tools", {}) all_tools: list[dict[str, Any]] = [] tool_by_name: dict[str, dict[str, Any]] = {} for name, defn in tools_data.items(): tool: dict[str, Any] = { "type": "function", "function": { "name": name, "description": defn.get("description", ""), "parameters": { "type": "object", "properties": defn.get("properties", {}), "required": defn.get("required", []), }, }, } all_tools.append(tool) tool_by_name[name] = tool scenarios: list[Scenario] = [] for s in data.get("scenarios", []): if "tools" in s: scenario_tools = [tool_by_name[t] for t in s["tools"]] else: scenario_tools = list(all_tools) messages: list[dict[str, Any]] = [] for msg in s.get("messages", []): m: dict[str, Any] = {"role": msg["role"]} if "content" in msg: m["content"] = msg["content"] if "tool_calls" in msg: m["tool_calls"] = [ { "id": tc["id"], "type": "function", "function": { "name": tc["name"], "arguments": json.dumps(tc["arguments"]), }, } for tc in msg["tool_calls"] ] if "tool_call_id" in msg: m["tool_call_id"] = msg["tool_call_id"] messages.append(m) tool_result: str | None = None if "tool_result" in s: tool_result = json.dumps(s["tool_result"]) scenarios.append( Scenario( name=s["name"], description=s["description"], messages=messages, tools=scenario_tools, expect_tool_call=s["expect_tool_call"], expected_function=s.get("expected_function"), required_arg_keys=s.get("required_arg_keys"), tool_result=tool_result, nested_array_key=s.get("nested_array_key"), required_item_keys=s.get("required_item_keys"), ) ) return scenarios ApiName = Literal["openai", "claude", "responses"] @dataclass class ParsedResponse: finish_reason: str # "tool_calls" | "stop" | ... has_tool_call: bool tool_call: dict[str, str] | None # {"id": ..., "name": ..., "arguments": ...} content: str | None @dataclass class ScenarioResult: name: str api: str phase: str # "tool_call" or "follow_up" passed: bool checks: dict[str, bool] = field(default_factory=dict) error: str | None = None latency_ms: float = 0.0 def validate_args(args_str: str, required_keys: list[str]) -> tuple[bool, str | None]: """Parse JSON arguments and check required keys exist.""" try: args = json.loads(args_str) except (json.JSONDecodeError, TypeError) as exc: return False, f"Invalid JSON: {exc}" if not isinstance(args, dict): return False, f"Expected dict, got {type(args).__name__}" missing = [k for k in required_keys if k not in args] if missing: return False, f"Missing keys: {missing}" return True, None def validate_nested_args( args_str: str, array_key: str, required_item_keys: list[str], ) -> tuple[bool, str | None]: """Check that args[array_key] is a list of objects with required keys.""" try: args = json.loads(args_str) except (json.JSONDecodeError, TypeError) as exc: return False, f"Invalid JSON: {exc}" if not isinstance(args, dict): return False, f"Expected dict, got {type(args).__name__}" arr = args.get(array_key) if not isinstance(arr, list): return False, f"'{array_key}' is not an array (got {type(arr).__name__})" if len(arr) == 0: return False, f"'{array_key}' is empty" for i, item in enumerate(arr): if not isinstance(item, dict): return ( False, f"'{array_key}[{i}]' is not an object (got {type(item).__name__})", ) missing = [k for k in required_item_keys if k not in item] if missing: return False, f"'{array_key}[{i}]' missing keys: {missing}" return True, None def call_api( client: httpx.Client, host: str, port: int, path: str, body: dict[str, Any], timeout: float, ) -> tuple[dict[str, Any], float]: """POST to http://{host}:{port}{path}, return (response_json, latency_ms).""" url = f"http://{host}:{port}{path}" t0 = time.monotonic() resp = client.post(url, json=body, timeout=timeout) latency = (time.monotonic() - t0) * 1000 resp.raise_for_status() return resp.json(), latency def _openai_build_request( model: str, messages: list[dict[str, Any]], tools: list[dict[str, Any]], ) -> tuple[str, dict[str, Any]]: """Build request for /v1/chat/completions.""" body: dict[str, Any] = { "model": model, "messages": messages, "tools": tools, "max_tokens": 16384, "temperature": 0.0, } return "/v1/chat/completions", body def _openai_parse_response(data: dict[str, Any]) -> ParsedResponse: """Parse OpenAI Chat Completions response into common format.""" choice = data["choices"][0] finish_reason = choice.get("finish_reason", "") message = choice.get("message", {}) tool_calls = message.get("tool_calls") content = message.get("content") has_tool_call = isinstance(tool_calls, list) and len(tool_calls) > 0 tool_call_info: dict[str, str] | None = None if has_tool_call: tc = tool_calls[0] fn = tc.get("function", {}) tool_call_info = { "id": tc.get("id", "call_0"), "name": fn.get("name", ""), "arguments": fn.get("arguments", "{}"), } return ParsedResponse( finish_reason=finish_reason, has_tool_call=has_tool_call, tool_call=tool_call_info, content=content, ) def _openai_build_followup( messages: list[dict[str, Any]], tools: list[dict[str, Any]], model: str, parsed: ParsedResponse, tool_result: str, ) -> tuple[str, dict[str, Any]]: """Build multi-turn follow-up for OpenAI Chat Completions.""" assert parsed.tool_call is not None tc = parsed.tool_call followup_messages: list[dict[str, Any]] = list(messages) + [ { "role": "assistant", "tool_calls": [ { "id": tc["id"], "type": "function", "function": { "name": tc["name"], "arguments": tc["arguments"], }, } ], }, { "role": "tool", "tool_call_id": tc["id"], "content": tool_result, }, ] body: dict[str, Any] = { "model": model, "messages": followup_messages, "tools": tools, "max_tokens": 16384, "temperature": 0.0, } return "/v1/chat/completions", body def _claude_translate_tools(tools: list[dict[str, Any]]) -> list[dict[str, Any]]: """Translate OpenAI-format tools to Claude format.""" claude_tools: list[dict[str, Any]] = [] for tool in tools: fn = tool["function"] claude_tools.append( { "name": fn["name"], "description": fn.get("description", ""), "input_schema": fn.get("parameters", {}), } ) return claude_tools def _claude_translate_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: """Translate OpenAI-format messages to Claude Messages format.""" claude_messages: list[dict[str, Any]] = [] for msg in messages: role = msg["role"] if role == "user": claude_messages.append( { "role": "user", "content": msg["content"], } ) elif role == "assistant": content_blocks: list[dict[str, Any]] = [] text_content = msg.get("content") if text_content and isinstance(text_content, str) and text_content.strip(): content_blocks.append({"type": "text", "text": text_content}) tool_calls = msg.get("tool_calls") if tool_calls: for tc in tool_calls: fn = tc.get("function", {}) args_str = fn.get("arguments", "{}") try: args_dict = json.loads(args_str) except (json.JSONDecodeError, TypeError): args_dict = {} content_blocks.append( { "type": "tool_use", "id": tc.get("id", "call_0"), "name": fn.get("name", ""), "input": args_dict, } ) if not content_blocks: content_blocks.append({"type": "text", "text": ""}) claude_messages.append( { "role": "assistant", "content": content_blocks, } ) elif role == "tool": claude_messages.append( { "role": "user", "content": [ { "type": "tool_result", "tool_use_id": msg.get("tool_call_id", "call_0"), "content": msg.get("content", ""), } ], } ) elif role == "system": pass return claude_messages def _claude_build_request( model: str, messages: list[dict[str, Any]], tools: list[dict[str, Any]], ) -> tuple[str, dict[str, Any]]: """Build request for /v1/messages.""" claude_messages = _claude_translate_messages(messages) claude_tools = _claude_translate_tools(tools) system_content: str | None = None for msg in messages: if msg["role"] == "system": system_content = msg["content"] break body: dict[str, Any] = { "model": model, "messages": claude_messages, "tools": claude_tools, "max_tokens": 16384, "temperature": 0.0, } if system_content is not None: body["system"] = system_content return "/v1/messages", body def _claude_parse_response(data: dict[str, Any]) -> ParsedResponse: """Parse Claude Messages response into common format.""" stop_reason = data.get("stop_reason", "") content_blocks = data.get("content", []) if stop_reason == "tool_use": finish_reason = "tool_calls" elif stop_reason == "end_turn": finish_reason = "stop" else: finish_reason = stop_reason tool_call_info: dict[str, str] | None = None text_parts: list[str] = [] has_tool_call = False for block in content_blocks: block_type = block.get("type") if block_type == "tool_use": has_tool_call = True if tool_call_info is None: input_data = block.get("input", {}) tool_call_info = { "id": block.get("id", "call_0"), "name": block.get("name", ""), "arguments": json.dumps(input_data) if isinstance(input_data, dict) else str(input_data), } elif block_type == "text": text = block.get("text", "") if text.strip(): text_parts.append(text) content = "\n".join(text_parts) if text_parts else None return ParsedResponse( finish_reason=finish_reason, has_tool_call=has_tool_call, tool_call=tool_call_info, content=content, ) def _claude_build_followup( messages: list[dict[str, Any]], tools: list[dict[str, Any]], model: str, parsed: ParsedResponse, tool_result: str, ) -> tuple[str, dict[str, Any]]: """Build multi-turn follow-up for Claude Messages.""" assert parsed.tool_call is not None tc = parsed.tool_call try: args_dict = json.loads(tc["arguments"]) except (json.JSONDecodeError, TypeError): args_dict = {} claude_messages = _claude_translate_messages(messages) claude_messages.append( { "role": "assistant", "content": [ { "type": "tool_use", "id": tc["id"], "name": tc["name"], "input": args_dict, } ], } ) claude_messages.append( { "role": "user", "content": [ { "type": "tool_result", "tool_use_id": tc["id"], "content": tool_result, } ], } ) claude_tools = _claude_translate_tools(tools) system_content: str | None = None for msg in messages: if msg["role"] == "system": system_content = msg["content"] break body: dict[str, Any] = { "model": model, "messages": claude_messages, "tools": claude_tools, "max_tokens": 16384, "temperature": 0.0, } if system_content is not None: body["system"] = system_content return "/v1/messages", body def _responses_translate_input(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: """Translate OpenAI chat messages to Responses API input items.""" items: list[dict[str, Any]] = [] for msg in messages: role = msg["role"] if role in ("user", "system"): items.append( { "type": "message", "role": role, "content": msg["content"], } ) elif role == "assistant": text_content = msg.get("content") if text_content and isinstance(text_content, str) and text_content.strip(): items.append( { "type": "message", "role": "assistant", "content": text_content, } ) tool_calls = msg.get("tool_calls") if tool_calls: for tc in tool_calls: fn = tc.get("function", {}) items.append( { "type": "function_call", "call_id": tc.get("id", "call_0"), "name": fn.get("name", ""), "arguments": fn.get("arguments", "{}"), } ) elif role == "tool": items.append( { "type": "function_call_output", "call_id": msg.get("tool_call_id", "call_0"), "output": msg.get("content", ""), } ) return items def _responses_build_request( model: str, messages: list[dict[str, Any]], tools: list[dict[str, Any]], ) -> tuple[str, dict[str, Any]]: """Build request for /v1/responses.""" input_items = _responses_translate_input(messages) body: dict[str, Any] = { "model": model, "input": input_items, "tools": tools, "temperature": 0.0, "max_output_tokens": 4096, } return "/v1/responses", body def _responses_parse_response(data: dict[str, Any]) -> ParsedResponse: """Parse OpenAI Responses API response into common format.""" output = data.get("output", []) tool_call_info: dict[str, str] | None = None text_parts: list[str] = [] has_tool_call = False for item in output: item_type = item.get("type") if item_type == "function_call": has_tool_call = True if tool_call_info is None: tool_call_info = { "id": item.get("call_id", "call_0"), "name": item.get("name", ""), "arguments": item.get("arguments", "{}"), } elif item_type == "message": msg_content = item.get("content", []) if isinstance(msg_content, list): for block in msg_content: if isinstance(block, dict): text = block.get("text", "") if text and text.strip(): text_parts.append(text) elif isinstance(msg_content, str) and msg_content.strip(): text_parts.append(msg_content) content = "\n".join(text_parts) if text_parts else None if has_tool_call: finish_reason = "tool_calls" else: status = data.get("status", "completed") finish_reason = "stop" if status == "completed" else status return ParsedResponse( finish_reason=finish_reason, has_tool_call=has_tool_call, tool_call=tool_call_info, content=content, ) def _responses_build_followup( messages: list[dict[str, Any]], tools: list[dict[str, Any]], model: str, parsed: ParsedResponse, tool_result: str, ) -> tuple[str, dict[str, Any]]: """Build multi-turn follow-up for Responses API.""" assert parsed.tool_call is not None tc = parsed.tool_call input_items = _responses_translate_input(messages) input_items.append( { "type": "function_call", "call_id": tc["id"], "name": tc["name"], "arguments": tc["arguments"], } ) input_items.append( { "type": "function_call_output", "call_id": tc["id"], "output": tool_result, } ) body: dict[str, Any] = { "model": model, "input": input_items, "tools": tools, "temperature": 0.0, "max_output_tokens": 4096, } return "/v1/responses", body ADAPTERS: dict[ApiName, dict[str, Any]] = { "openai": { "build_request": _openai_build_request, "parse_response": _openai_parse_response, "build_followup": _openai_build_followup, }, "claude": { "build_request": _claude_build_request, "parse_response": _claude_parse_response, "build_followup": _claude_build_followup, }, "responses": { "build_request": _responses_build_request, "parse_response": _responses_parse_response, "build_followup": _responses_build_followup, }, } def run_scenario( client: httpx.Client, host: str, port: int, model: str, scenario: Scenario, api_name: ApiName, timeout: float, verbose: bool, ) -> list[ScenarioResult]: """Run a single scenario against one API adapter. Returns 1-2 results.""" adapter = ADAPTERS[api_name] build_request = adapter["build_request"] parse_response = adapter["parse_response"] build_followup = adapter["build_followup"] results: list[ScenarioResult] = [] # --- Phase 1: initial request --- path, body = build_request(model, scenario.messages, scenario.tools) if verbose: print( f" [{api_name}] request: {path} {json.dumps(body, indent=2)}", file=sys.stderr, ) try: data, latency = call_api(client, host, port, path, body, timeout) except Exception as exc: results.append( ScenarioResult( name=scenario.name, api=api_name, phase="tool_call", passed=False, error=f"API error: {exc}", ) ) return results if verbose: print( f" [{api_name}] response: {json.dumps(data, indent=2)}", file=sys.stderr ) parsed = parse_response(data) checks: dict[str, bool] = {} if scenario.expect_tool_call: checks["finish_reason_tool_calls"] = parsed.finish_reason == "tool_calls" checks["has_tool_call"] = parsed.has_tool_call args_err: str | None = None if parsed.has_tool_call and parsed.tool_call is not None: checks["correct_function"] = ( scenario.expected_function is None or parsed.tool_call["name"] == scenario.expected_function ) if scenario.required_arg_keys: ok, args_err = validate_args( parsed.tool_call["arguments"], scenario.required_arg_keys ) checks["valid_arguments"] = ok else: checks["valid_arguments"] = True if scenario.nested_array_key and scenario.required_item_keys: ok, nested_err = validate_nested_args( parsed.tool_call["arguments"], scenario.nested_array_key, scenario.required_item_keys, ) checks["valid_nested_structure"] = ok if not ok: args_err = nested_err else: checks["correct_function"] = False checks["valid_arguments"] = False args_err = "No tool call returned" passed = all(checks.values()) error = args_err if not passed else None else: checks["finish_reason_stop"] = parsed.finish_reason == "stop" checks["no_tool_call"] = not parsed.has_tool_call checks["has_content"] = ( parsed.content is not None and len(parsed.content.strip()) > 0 ) passed = all(checks.values()) error = ( None if passed else ( f"finish_reason={parsed.finish_reason}, " f"tool_call={'yes' if parsed.has_tool_call else 'no'}, " f"content={'yes' if parsed.content else 'no'}" ) ) results.append( ScenarioResult( name=scenario.name, api=api_name, phase="tool_call", passed=passed, checks=checks, error=error, latency_ms=latency, ) ) # --- Phase 2: multi-turn follow-up --- if ( scenario.tool_result is not None and parsed.has_tool_call and parsed.tool_call is not None ): followup_path, followup_body = build_followup( scenario.messages, scenario.tools, model, parsed, scenario.tool_result, ) if verbose: print( f" [{api_name}] follow_up request: {followup_path} {json.dumps(followup_body, indent=2)}", file=sys.stderr, ) try: data2, latency2 = call_api( client, host, port, followup_path, followup_body, timeout ) except Exception as exc: results.append( ScenarioResult( name=scenario.name, api=api_name, phase="follow_up", passed=False, error=f"API error: {exc}", ) ) return results if verbose: print( f" [{api_name}] follow_up response: {json.dumps(data2, indent=2)}", file=sys.stderr, ) parsed2 = parse_response(data2) checks2: dict[str, bool] = {} checks2["finish_reason_stop"] = parsed2.finish_reason == "stop" checks2["no_tool_call"] = not parsed2.has_tool_call checks2["has_content"] = ( parsed2.content is not None and len(parsed2.content.strip()) > 0 ) passed2 = all(checks2.values()) error2: str | None = None if not passed2: error2 = ( f"finish_reason={parsed2.finish_reason}, " f"tool_call={'yes' if parsed2.has_tool_call else 'no'}, " f"content={'yes' if parsed2.content else 'no'}" ) results.append( ScenarioResult( name=scenario.name, api=api_name, phase="follow_up", passed=passed2, checks=checks2, error=error2, latency_ms=latency2, ) ) return results def result_to_dict(result: ScenarioResult) -> dict[str, Any]: """Convert a ScenarioResult to a JSON-serializable dict.""" return { "name": result.name, "api": result.api, "phase": result.phase, "passed": result.passed, "checks": result.checks, "error": result.error, "latency_ms": round(result.latency_ms, 1), } _MULTI_NODE_PRIORITY: dict[tuple[str, str], int] = { ("tensor", "jaccl"): 0, ("pipeline", "jaccl"): 2, ("pipeline", "ring"): 3, ("tensor", "ring"): 4, } _SINGLE_NODE_PRIORITY = 1 def _placement_sort_key(p: dict[str, Any]) -> tuple[int, int]: sharding = p.get("sharding", "").lower() meta = p.get("instance_meta", "").lower() kind = ( "tensor" if "tensor" in sharding else "pipeline", "jaccl" if "jaccl" in meta else "ring", ) n_nodes = nodes_used_in_instance(p["instance"]) if n_nodes == 1: return (_SINGLE_NODE_PRIORITY, -n_nodes) priority = _MULTI_NODE_PRIORITY.get(kind, 99) return (priority, -n_nodes) def main() -> None: parser = argparse.ArgumentParser( description="Multi-API tool-calling eval for exo", formatter_class=argparse.RawDescriptionHelpFormatter, epilog="""\ Examples: %(prog)s --model mlx-community/Qwen3-30B-A3B-4bit %(prog)s --model my-model --api openai --repeat 3 %(prog)s --model my-model --api all --scenarios weather_simple calculator_multi_turn %(prog)s --model my-model --stdout """, ) add_common_instance_args(parser) parser.add_argument( "--api", choices=["openai", "claude", "responses", "all"], default="all", help="Which API adapter(s) to test (default: all)", ) parser.add_argument( "--repeat", type=int, default=1, help="Repeat each scenario N times (default: 1)", ) parser.add_argument( "--scenarios", nargs="*", help="Run only these scenarios (by name)", ) parser.add_argument( "--verbose", action="store_true", help="Print full API responses to stderr", ) parser.add_argument( "--json-out", default="bench/eval_results.json", help="Write JSON results to file (default: bench/eval_results.json)", ) parser.add_argument( "--stdout", action="store_true", help="Write JSON results to stdout instead of file", ) args = parser.parse_args() all_scenarios = load_scenarios(SCENARIOS_PATH) if args.scenarios: scenarios = [s for s in all_scenarios if s.name in args.scenarios] if not scenarios: print( f"No matching scenarios. Available: {[s.name for s in all_scenarios]}", file=sys.stderr, ) sys.exit(1) else: scenarios = all_scenarios api_names: list[ApiName] = ( ["openai", "claude", "responses"] if args.api == "all" else [args.api] ) log = sys.stderr if args.stdout else sys.stdout exo = ExoClient(args.host, args.port, timeout_s=args.timeout) _short_id, full_model_id = resolve_model_short_id(exo, args.model) selected = settle_and_fetch_placements( exo, full_model_id, args, settle_timeout=args.settle_timeout ) if not selected: print("No valid placements matched your filters.", file=sys.stderr) sys.exit(1) selected.sort(key=_placement_sort_key) preview = selected[0] settle_deadline = ( time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None ) print("Planning phase: checking downloads...", file=log) run_planning_phase( exo, full_model_id, preview, args.danger_delete_downloads, args.timeout, settle_deadline, ) instance = preview["instance"] instance_id = instance_id_from_instance(instance) sharding = str(preview["sharding"]) instance_meta = str(preview["instance_meta"]) n_nodes = nodes_used_in_instance(instance) print(f"Model: {full_model_id}", file=log) print(f"Placement: {sharding} / {instance_meta} / {n_nodes} nodes", file=log) print(f"Endpoint: http://{args.host}:{args.port}", file=log) print(f"APIs: {', '.join(api_names)}", file=log) total_runs = len(scenarios) * args.repeat * len(api_names) print( f"Scenarios: {len(scenarios)} x {args.repeat} repeats x {len(api_names)} APIs = {total_runs} runs", file=log, ) print("=" * 72, file=log) exo.request_json("POST", "/instance", body={"instance": instance}) try: wait_for_instance_ready(exo, instance_id) except (RuntimeError, TimeoutError) as e: print(f"Failed to initialize placement: {e}", file=sys.stderr) with contextlib.suppress(ExoHttpError): exo.request_json("DELETE", f"/instance/{instance_id}") sys.exit(1) time.sleep(1) all_results: list[ScenarioResult] = [] try: with httpx.Client() as http_client: for run_idx in range(args.repeat): if args.repeat > 1: print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log) for scenario in scenarios: for api_name in api_names: print( f"\n [{api_name:>9}] {scenario.name}: {scenario.description}", file=log, ) scenario_results = run_scenario( http_client, args.host, args.port, full_model_id, scenario, api_name, args.timeout, args.verbose, ) all_results.extend(scenario_results) for r in scenario_results: status = "PASS" if r.passed else "FAIL" print( f" [{r.phase:>10}] {status} ({r.latency_ms:.0f}ms)", file=log, ) for check_name, check_ok in r.checks.items(): mark = "+" if check_ok else "-" print(f" {mark} {check_name}", file=log) if r.error: print(f" ! {r.error}", file=log) finally: try: exo.request_json("DELETE", f"/instance/{instance_id}") except ExoHttpError as e: if e.status != 404: raise wait_for_instance_gone(exo, instance_id) # --- Summary --- print(f"\n{'=' * 72}", file=log) total = len(all_results) passed = sum(1 for r in all_results if r.passed) tool_call_results = [r for r in all_results if r.phase == "tool_call"] follow_up_results = [r for r in all_results if r.phase == "follow_up"] tc_passed = sum(1 for r in tool_call_results if r.passed) fu_passed = sum(1 for r in follow_up_results if r.passed) avg_latency = sum(r.latency_ms for r in all_results) / total if total else 0 print( f"Total: {passed}/{total} passed ({100 * passed / total:.0f}%)", file=log ) print(f"Tool call: {tc_passed}/{len(tool_call_results)} passed", file=log) if follow_up_results: print(f"Follow-up: {fu_passed}/{len(follow_up_results)} passed", file=log) print(f"Avg latency: {avg_latency:.0f}ms", file=log) for api_name in api_names: api_results = [r for r in all_results if r.api == api_name] api_passed = sum(1 for r in api_results if r.passed) print(f" {api_name:>9}: {api_passed}/{len(api_results)} passed", file=log) if passed < total: print("\nFailed:", file=log) for r in all_results: if not r.passed: print(f" - {r.name} [{r.api}/{r.phase}]: {r.error}", file=log) json_results = [result_to_dict(r) for r in all_results] if args.stdout: print(json.dumps(json_results, indent=2)) else: json_path = args.json_out parent = os.path.dirname(json_path) if parent: os.makedirs(parent, exist_ok=True) with open(json_path, "w") as f: json.dump(json_results, f, indent=2) f.write("\n") print(f"\nJSON results written to {json_path}", file=log) sys.exit(0 if passed == total else 1) if __name__ == "__main__": main() ================================================ FILE: bench/exo_bench.py ================================================ # type: ignore #!/usr/bin/env python3 """Tool-calling eval for exo's OpenAI-compatible API. Tests whether models correctly: - Trigger tool calls when appropriate - Return valid JSON arguments matching function schemas - Handle multi-turn tool use (call -> result -> final answer) - Avoid calling tools when unnecessary Start exo with a model first, then run: uv run python tool_call_eval.py --model uv run python tool_call_eval.py --model --host 10.0.0.5 --port 52415 uv run python tool_call_eval.py --model --repeat 3 uv run python tool_call_eval.py --model --scenarios weather_simple calculator_multi_turn """ from __future__ import annotations import argparse import contextlib import itertools import json import sys import time from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path from statistics import mean from typing import Any from harness import ( ExoClient, ExoHttpError, add_common_instance_args, instance_id_from_instance, nodes_used_in_instance, resolve_model_short_id, run_planning_phase, settle_and_fetch_placements, wait_for_instance_gone, wait_for_instance_ready, ) from loguru import logger from transformers import AutoTokenizer # Monkey-patch for transformers 5.x compatibility # Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location # which was moved in transformers 5.0.0rc2 try: import transformers.models.gpt2.tokenization_gpt2 as gpt2_tokenization from transformers.convert_slow_tokenizer import bytes_to_unicode if not hasattr(gpt2_tokenization, "bytes_to_unicode"): gpt2_tokenization.bytes_to_unicode = bytes_to_unicode # type: ignore[attr-defined] except ImportError: pass # transformers < 5.0 or bytes_to_unicode not available def load_tokenizer_for_bench(model_id: str) -> Any: """ Load tokenizer for benchmarking, with special handling for Kimi models. Kimi uses a custom TikTokenTokenizer that transformers 5.x can't load via AutoTokenizer. This function replicates the logic from utils_mlx.py for bench compatibility. """ model_id_lower = model_id.lower() if "kimi-k2" in model_id_lower: import importlib.util import types from huggingface_hub import snapshot_download # Download/get the model path model_path = Path( snapshot_download( model_id, allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model"], ) ) sys.path.insert(0, str(model_path)) # Load tool_declaration_ts first (tokenization_kimi imports it with relative import) tool_decl_path = model_path / "tool_declaration_ts.py" if tool_decl_path.exists(): spec = importlib.util.spec_from_file_location( "tool_declaration_ts", tool_decl_path ) if spec and spec.loader: tool_decl_module = importlib.util.module_from_spec(spec) sys.modules["tool_declaration_ts"] = tool_decl_module spec.loader.exec_module(tool_decl_module) # Load tokenization_kimi with patched source (convert relative to absolute import) tok_path = model_path / "tokenization_kimi.py" source = tok_path.read_text() source = source.replace("from .tool_declaration_ts", "from tool_declaration_ts") spec = importlib.util.spec_from_file_location("tokenization_kimi", tok_path) if spec: tok_module = types.ModuleType("tokenization_kimi") tok_module.__file__ = str(tok_path) sys.modules["tokenization_kimi"] = tok_module exec(compile(source, tok_path, "exec"), tok_module.__dict__) # noqa: S102 TikTokenTokenizer = tok_module.TikTokenTokenizer # noqa: N806 else: from tokenization_kimi import TikTokenTokenizer # type: ignore[import-not-found] # noqa: I001 hf_tokenizer: Any = TikTokenTokenizer.from_pretrained(model_path) # Patch encode to use internal tiktoken model directly # transformers 5.x has a bug in the encode->pad path for slow tokenizers def _patched_encode(text: str, **kwargs: object) -> list[int]: # Pass allowed_special="all" to handle special tokens like <|im_user|> return list(hf_tokenizer.model.encode(text, allowed_special="all")) hf_tokenizer.encode = _patched_encode return hf_tokenizer # Default: use AutoTokenizer return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) def format_peak_memory(b: float) -> str: for unit in ["B", "KB", "MB", "GB", "TB"]: if b < 1024.0: return f"{b:.2f}{unit}" b /= 1024.0 raise ValueError("You're using petabytes of memory. Something went wrong...") def parse_int_list(values: list[str]) -> list[int]: items: list[int] = [] for v in values: for part in v.split(","): part = part.strip() if part: items.append(int(part)) return items def run_one_completion( client: ExoClient, model_id: str, pp_hint: int, tg: int, prompt_sizer: PromptSizer ) -> tuple[dict[str, Any], int]: content, pp_tokens = prompt_sizer.build(pp_hint) payload: dict[str, Any] = { "model": model_id, "messages": [{"role": "user", "content": content}], "stream": False, "max_tokens": tg, } t0 = time.perf_counter() out = client.post_bench_chat_completions(payload) elapsed = time.perf_counter() - t0 stats = out.get("generation_stats") # Extract preview, handling None content (common for thinking models) choices = out.get("choices") or [{}] message = choices[0].get("message", {}) if choices else {} content = message.get("content") or "" preview = content[:200] if content else "" return { "elapsed_s": elapsed, "output_text_preview": preview, "stats": stats, }, pp_tokens class PromptSizer: def __init__(self, tokenizer: Any, atom: str = "a "): self.tokenizer = tokenizer self.atom = atom self.count_fn = PromptSizer._make_counter(tokenizer) self.base_tokens = self.count_fn("") @staticmethod def _make_counter(tokenizer: Any) -> Callable[[str], int]: def count_fn(user_content: str) -> int: messages = [{"role": "user", "content": user_content}] ids = tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=True ) # Fix for transformers 5.x if hasattr(ids, "input_ids"): ids = ids.input_ids return int(len(ids)) return count_fn def build(self, target_prompt_tokens: int) -> tuple[str, int]: target = int(target_prompt_tokens) if target < self.base_tokens: raise RuntimeError( f"Target ({target}) is smaller than template overhead ({self.base_tokens})." ) # Estimate tokens per atom using a sample sample_count = 100 sample_content = self.atom * sample_count sample_tokens = self.count_fn(sample_content) - self.base_tokens tokens_per_atom = sample_tokens / sample_count # Estimate starting point needed_tokens = target - self.base_tokens estimated_atoms = int(needed_tokens / tokens_per_atom) # Binary search to find exact atom count low, high = 0, estimated_atoms * 2 + 100 while low < high: mid = (low + high) // 2 tok = self.count_fn(self.atom * mid) if tok < target: low = mid + 1 else: high = mid content = self.atom * low tok = self.count_fn(content) logger.info(f"{tok=}") if tok != target: raise RuntimeError( f"Overshot: got {tok} tokens (target {target}). " f"Pick a different atom (try ' a' or '\\n' or '0 ')." ) return content, tok def main() -> int: ap = argparse.ArgumentParser( prog="exo-bench", description="Benchmark exo model throughput across placement previews.", ) add_common_instance_args(ap) ap.add_argument( "--pp", nargs="+", required=True, help="Prompt-size hints (ints). Accepts commas.", ) ap.add_argument( "--tg", nargs="+", required=True, help="Generation lengths (ints). Accepts commas.", ) ap.add_argument( "--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair." ) ap.add_argument( "--concurrency", nargs="+", default=["1"], help="Concurrency levels (ints). Accepts commas. E.g. --concurrency 1,2,4,8. Default 1.", ) ap.add_argument( "--warmup", type=int, default=0, help="Warmup runs per placement (uses first pp/tg).", ) ap.add_argument( "--json-out", default="bench/results.json", help="Write raw per-run results JSON to this path.", ) ap.add_argument("--stdout", action="store_true", help="Write results to stdout") ap.add_argument( "--dry-run", action="store_true", help="List selected placements and exit." ) ap.add_argument( "--all-combinations", action="store_true", help="Force all pp×tg combinations (cartesian product) even when lists have equal length.", ) args = ap.parse_args() pp_list = parse_int_list(args.pp) tg_list = parse_int_list(args.tg) if not pp_list or not tg_list: logger.error("pp and tg lists must be non-empty") return 2 if args.repeat <= 0: logger.error("--repeat must be >= 1") return 2 concurrency_list = parse_int_list(args.concurrency) if not concurrency_list or any(c <= 0 for c in concurrency_list): logger.error("--concurrency values must be >= 1") return 2 # Log pairing mode use_combinations = args.all_combinations or len(pp_list) != len(tg_list) if use_combinations: logger.info( f"pp/tg mode: combinations (product) - {len(pp_list) * len(tg_list)} pairs" ) else: logger.info(f"pp/tg mode: tandem (zip) - {len(pp_list)} pairs") client = ExoClient(args.host, args.port, timeout_s=args.timeout) short_id, full_model_id = resolve_model_short_id( client, args.model, force_download=args.force_download ) tokenizer = load_tokenizer_for_bench(full_model_id) if tokenizer is None: raise RuntimeError("[exo-bench] tokenizer load failed") try: prompt_sizer = PromptSizer(tokenizer) logger.debug(f"[exo-bench] loaded tokenizer: {full_model_id} for prompt sizer") except Exception: logger.error("[exo-bench] tokenizer usable but prompt sizing failed") raise selected = settle_and_fetch_placements( client, full_model_id, args, settle_timeout=args.settle_timeout ) if not selected: logger.error("No valid placements matched your filters.") return 1 selected.sort( key=lambda p: ( str(p.get("instance_meta", "")), str(p.get("sharding", "")), -nodes_used_in_instance(p["instance"]), ), reverse=True, ) logger.debug(f"exo-bench model: short_id={short_id} full_id={full_model_id}") logger.info(f"placements: {len(selected)}") for p in selected: logger.info( f" - {p['sharding']} / {p['instance_meta']} / nodes={nodes_used_in_instance(p['instance'])}" ) if args.dry_run: return 0 settle_deadline = ( time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None ) logger.info("Planning phase: checking downloads...") download_duration_s = run_planning_phase( client, full_model_id, selected[0], args.danger_delete_downloads, args.timeout, settle_deadline, ) if download_duration_s is not None: logger.info(f"Download: {download_duration_s:.1f}s (freshly downloaded)") else: logger.info("Download: model already cached") all_rows: list[dict[str, Any]] = [] for preview in selected: instance = preview["instance"] instance_id = instance_id_from_instance(instance) sharding = str(preview["sharding"]) instance_meta = str(preview["instance_meta"]) n_nodes = nodes_used_in_instance(instance) logger.info("=" * 80) logger.info( f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes} / instance_id={instance_id}" ) client.request_json("POST", "/instance", body={"instance": instance}) try: wait_for_instance_ready(client, instance_id) except (RuntimeError, TimeoutError) as e: logger.error(f"Failed to initialize placement: {e}") with contextlib.suppress(ExoHttpError): client.request_json("DELETE", f"/instance/{instance_id}") continue time.sleep(1) try: for i in range(args.warmup): run_one_completion( client, full_model_id, pp_list[0], tg_list[0], prompt_sizer ) logger.debug(f" warmup {i + 1}/{args.warmup} done") # If pp and tg lists have same length, run in tandem (zip) # Otherwise (or if --all-combinations), run all combinations (cartesian product) if use_combinations: pp_tg_pairs = list(itertools.product(pp_list, tg_list)) else: pp_tg_pairs = list(zip(pp_list, tg_list, strict=True)) for pp, tg in pp_tg_pairs: for concurrency in concurrency_list: logger.info(f"--- pp={pp} tg={tg} concurrency={concurrency} ---") runs: list[dict[str, Any]] = [] for r in range(args.repeat): time.sleep(3) if concurrency <= 1: # Sequential: single request try: row, actual_pp_tokens = run_one_completion( client, full_model_id, pp, tg, prompt_sizer ) except Exception as e: logger.error(e) continue row.update( { "model_short_id": short_id, "model_id": full_model_id, "placement_sharding": sharding, "placement_instance_meta": instance_meta, "placement_nodes": n_nodes, "instance_id": instance_id, "pp_tokens": actual_pp_tokens, "tg": tg, "repeat_index": r, "concurrency": 1, **( {"download_duration_s": download_duration_s} if download_duration_s is not None else {} ), } ) runs.append(row) all_rows.append(row) else: # Concurrent: fire N requests in parallel # Each thread gets its own ExoClient (separate HTTP connection) batch_results: list[tuple[dict[str, Any], int]] = [] batch_errors = 0 def _run_concurrent( idx: int, *, _pp: int = pp, _tg: int = tg ) -> tuple[dict[str, Any], int]: c = ExoClient( args.host, args.port, timeout_s=args.timeout ) return run_one_completion( c, full_model_id, _pp, _tg, prompt_sizer ) with ThreadPoolExecutor(max_workers=concurrency) as pool: futures = { pool.submit(_run_concurrent, i): i for i in range(concurrency) } for fut in as_completed(futures): try: batch_results.append(fut.result()) except Exception as e: logger.error(f"Concurrent request failed: {e}") batch_errors += 1 for idx, (row, actual_pp_tokens) in enumerate( batch_results ): row.update( { "model_short_id": short_id, "model_id": full_model_id, "placement_sharding": sharding, "placement_instance_meta": instance_meta, "placement_nodes": n_nodes, "instance_id": instance_id, "pp_tokens": actual_pp_tokens, "tg": tg, "repeat_index": r, "concurrency": concurrency, "concurrent_index": idx, **( {"download_duration_s": download_duration_s} if download_duration_s is not None else {} ), } ) runs.append(row) all_rows.append(row) if batch_results: valid_gen_tps = [ x["stats"]["generation_tps"] for x, _ in batch_results if x["stats"]["generation_tps"] > 0 ] agg_gen_tps = ( mean(valid_gen_tps) if valid_gen_tps else 0.0 ) gen_tps = agg_gen_tps / concurrency logger.info( f"[concurrent {concurrency}x] " f"agg_gen_tps={agg_gen_tps:.2f} " f"gen_tps={gen_tps:.2f} " f"errors={batch_errors}" ) if runs: prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs) gen_tps = mean( x["stats"]["generation_tps"] / x["concurrency"] for x in runs ) ptok = mean(x["stats"]["prompt_tokens"] for x in runs) gtok = mean(x["stats"]["generation_tokens"] for x in runs) peak = mean( x["stats"]["peak_memory_usage"]["inBytes"] for x in runs ) logger.info( f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f} " f"prompt_tokens={ptok} gen_tokens={gtok} " f"peak_memory={format_peak_memory(peak)}\n" ) time.sleep(2) finally: try: client.request_json("DELETE", f"/instance/{instance_id}") except ExoHttpError as e: if e.status != 404: raise wait_for_instance_gone(client, instance_id) logger.debug(f"Deleted instance {instance_id}") time.sleep(5) if args.stdout: json.dump(all_rows, sys.stdout, indent=2, ensure_ascii=False) elif args.json_out: with open(args.json_out, "w", encoding="utf-8") as f: json.dump(all_rows, f, indent=2, ensure_ascii=False) logger.debug(f"\nWrote results JSON: {args.json_out}") return 0 if __name__ == "__main__": raise SystemExit(main()) ================================================ FILE: bench/exo_eval.py ================================================ # type: ignore #!/usr/bin/env python3 """Quality evaluation for exo — matches Artificial Analysis methodology. Runs LLM benchmarks against exo's OpenAI-compatible API using the same prompts, temperature settings, and answer extraction as Artificial Analysis. Supported benchmarks: gpqa_diamond - Graduate-level science QA (198 questions, 4-choice MC) mmlu_pro - Multi-task language understanding (12K questions, 10-choice MC) aime_2024 - Math olympiad 2024 (30 problems, integer answers) aime_2025 - Math olympiad 2025 (30 problems, integer answers) humaneval - Python code generation (164 problems, pass@1) livecodebench - Competitive programming (880+ problems, pass@1) Model configs in eval_configs/models.toml auto-detect reasoning/non-reasoning settings per model. Override with --reasoning / --no-reasoning. Usage: uv run python exo_eval.py --model --tasks gpqa_diamond uv run python exo_eval.py --model --tasks humaneval,livecodebench --limit 50 uv run python exo_eval.py --model --tasks gpqa_diamond --compare-concurrency 1,4 References: https://artificialanalysis.ai/methodology/intelligence-benchmarking """ from __future__ import annotations import argparse import asyncio import contextlib import json import multiprocessing import random import re import sys import time import tomllib from dataclasses import dataclass from pathlib import Path from typing import Any import httpx from harness import ( ExoClient, ExoHttpError, add_common_instance_args, instance_id_from_instance, nodes_used_in_instance, resolve_model_short_id, run_planning_phase, settle_and_fetch_placements, wait_for_instance_gone, wait_for_instance_ready, ) from loguru import logger # --------------------------------------------------------------------------- # Artificial Analysis constants # --------------------------------------------------------------------------- MAX_RETRIES = 30 DEFAULT_MAX_TOKENS = 16_384 REASONING_MAX_TOKENS = 131_072 TEMPERATURE_NON_REASONING = 0.0 TEMPERATURE_REASONING = 1.0 # MC answer extraction: 8 fallback regex patterns. # All patterns are tried; the match at the latest text position wins # (handles models that self-correct during reasoning). _MC_PATTERNS: list[re.Pattern[str]] = [ re.compile( r"(?i)[\*\_]{0,2}Answer[\*\_]{0,2}\s*:[\s\*\_]{0,2}\s*([A-Z])(?![a-zA-Z0-9])" ), re.compile(r"\\boxed\{[^}]*([A-Z])[^}]*\}"), re.compile(r"(?i)answer is ([a-zA-Z])"), re.compile(r"(?i)answer is \\\(([a-zA-Z])"), re.compile(r"([A-Z])\)\s*[^A-Z]*$"), re.compile(r"([A-Z])\s+is\s+the\s+correct\s+answer"), re.compile(r"([A-Z])\s*$"), re.compile(r"([A-Z])\s*\."), ] # Code extraction: last ```python ... ``` block (AA regex) _CODE_BLOCK_RE = re.compile(r"```(?:python|Python)?\s*\n(.*?)```", re.DOTALL) # --------------------------------------------------------------------------- # Model config loading # --------------------------------------------------------------------------- def load_model_config(model_id: str) -> dict[str, Any] | None: """Look up model in eval_configs/models.toml. Returns config dict or None.""" config_path = Path(__file__).resolve().parent / "eval_configs" / "models.toml" if not config_path.exists(): return None with open(config_path, "rb") as f: data = tomllib.load(f) for entry in data.get("model", []): patterns = entry.get("patterns", []) if any(p in model_id for p in patterns): return entry return None # --------------------------------------------------------------------------- # Answer extraction # --------------------------------------------------------------------------- def extract_mc_answer(text: str, valid_letters: str = "ABCD") -> str | None: """Extract MC answer. Last match by text position wins.""" valid_set = set(valid_letters) best: tuple[int, str] | None = None for pattern in _MC_PATTERNS: for m in pattern.finditer(text): letter = m.group(1).upper() if letter in valid_set: pos = m.start() if best is None or pos >= best[0]: best = (pos, letter) return best[1] if best else None def extract_boxed_answer(text: str) -> str | None: r"""Extract content from the last \boxed{...}.""" matches: list[str] = [] idx = 0 while True: pos = text.find("\\boxed{", idx) if pos < 0: break depth = 0 i = pos + len("\\boxed{") start = i while i < len(text): if text[i] == "{": depth += 1 elif text[i] == "}": if depth == 0: matches.append(text[start:i]) break depth -= 1 i += 1 idx = i + 1 if i < len(text) else len(text) return matches[-1].strip() if matches else None def extract_code_block(text: str, preserve_indent: bool = False) -> str | None: """Extract the last Python code block from markdown response. If preserve_indent is True, only strip trailing whitespace (keeps leading indentation intact — needed for HumanEval function-body completions). """ matches = _CODE_BLOCK_RE.findall(text) if matches: raw = matches[-1] return raw.rstrip() if preserve_indent else raw.strip() # Fallback: try raw code after last ``` lines = text.split("\n") backtick_lines = [i for i, line in enumerate(lines) if "```" in line] if len(backtick_lines) >= 2: return "\n".join(lines[backtick_lines[-2] + 1 : backtick_lines[-1]]) return None def check_aime_answer(extracted: str, gold: int) -> bool: """Check if extracted AIME answer matches gold integer.""" try: return int(extracted.strip()) == gold except ValueError: pass try: from math_verify import parse, verify return verify(parse(str(gold)), parse(extracted)) except Exception: return False # --------------------------------------------------------------------------- # Code execution — official evaluation harnesses # --------------------------------------------------------------------------- # LiveCodeBench: vendored from https://github.com/LiveCodeBench/LiveCodeBench # run_test() must execute in a child process because reliability_guard() # permanently disables OS functions (os.kill, subprocess.Popen, etc.). def _lcb_worker( sample: dict, code: str, timeout: int, result_holder: list[Any], metadata_holder: list[Any], ) -> None: """Target for multiprocessing.Process — runs vendored LCB run_test.""" from vendor.lcb_testing_util import run_test try: results, metadata = run_test(sample, test=code, debug=False, timeout=timeout) result_holder.append(results) metadata_holder.append(metadata) except Exception as e: result_holder.append([-4]) metadata_holder.append({"error_code": -4, "error_message": str(e)}) def run_livecodebench_test( code: str, sample: dict, timeout: int = 6, ) -> tuple[bool, str]: """Run LCB evaluation in a subprocess. Returns (passed, diagnostic_info).""" manager = multiprocessing.Manager() result_holder = manager.list() metadata_holder = manager.list() proc = multiprocessing.Process( target=_lcb_worker, args=(sample, code, timeout, result_holder, metadata_holder), ) proc.start() # Global timeout: (per-test timeout + 1) * num_tests + 5 num_tests = len(json.loads(sample["input_output"]).get("inputs", [])) global_timeout = (timeout + 1) * num_tests + 5 proc.join(timeout=global_timeout) if proc.is_alive(): proc.kill() proc.join() return False, "Global timeout exceeded" if not result_holder: return False, "No results returned from worker" results = list(result_holder[0]) metadata = dict(metadata_holder[0]) if metadata_holder else {} # LCB convention: True = pass, negative int = failure code all_passed = all(r is True or r == 1 for r in results) if all_passed: return True, "" diag = metadata.get("error_message", "") if not diag and "output" in metadata: diag = f"Got {metadata['output']}, expected {metadata.get('expected', '?')}" return False, diag def run_humaneval_test( problem: dict, completion: str, timeout: float = 10.0 ) -> tuple[bool, str]: """Run HumanEval evaluation using the official human_eval package.""" from human_eval.execution import check_correctness result = check_correctness(problem, completion, timeout) passed = result["passed"] diag = "" if passed else result.get("result", "failed") return passed, diag # --------------------------------------------------------------------------- # Benchmark definitions # --------------------------------------------------------------------------- @dataclass class QuestionResult: question_id: int prompt: str response: str extracted_answer: str | None gold_answer: str correct: bool error: str | None = None prompt_tokens: int = 0 completion_tokens: int = 0 reasoning_tokens: int = 0 elapsed_s: float = 0.0 @dataclass class BenchmarkConfig: name: str description: str dataset_name: str dataset_config: str | None split: str kind: str # "mc", "math", "code" BENCHMARKS: dict[str, BenchmarkConfig] = { "gpqa_diamond": BenchmarkConfig( name="gpqa_diamond", description="Graduate-level science QA (198 Q, 4-choice MC)", dataset_name="Idavidrein/gpqa", dataset_config="gpqa_diamond", split="train", kind="mc", ), "mmlu_pro": BenchmarkConfig( name="mmlu_pro", description="Multi-task language understanding (12K Q, 10-choice MC)", dataset_name="TIGER-Lab/MMLU-Pro", dataset_config=None, split="test", kind="mc", ), "aime_2024": BenchmarkConfig( name="aime_2024", description="Math olympiad 2024 (30 problems, integer answers)", dataset_name="HuggingFaceH4/aime_2024", dataset_config=None, split="train", kind="math", ), "aime_2025": BenchmarkConfig( name="aime_2025", description="Math olympiad 2025 (30 problems, integer answers)", dataset_name="MathArena/aime_2025", dataset_config=None, split="train", kind="math", ), "humaneval": BenchmarkConfig( name="humaneval", description="Python code generation (164 problems, pass@1)", dataset_name="openai/openai_humaneval", dataset_config=None, split="test", kind="code", ), "livecodebench": BenchmarkConfig( name="livecodebench", description="Competitive programming (880+ problems, pass@1)", dataset_name="livecodebench/code_generation_lite", dataset_config=None, split="test", kind="code", ), } # --------------------------------------------------------------------------- # Prompt formatters # --------------------------------------------------------------------------- _GPQA_INSTRUCTION = ( "Answer the following multiple choice question. " "The last line of your response should be in the following format: " "'Answer: A/B/C/D' (e.g. 'Answer: A')." ) _MMLU_PRO_INSTRUCTION = ( "Answer the following multiple choice question. " "The last line of your response should be in the following format: " "'Answer: A/B/C/D/E/F/G/H/I/J' (e.g. 'Answer: A')." ) _AIME_INSTRUCTION = ( "Solve the following math problem step by step. " "Put your answer inside \\boxed{}.\n" "Remember to put your answer inside \\boxed{}." ) _HUMANEVAL_INSTRUCTION = ( "Complete the following Python function. Return only the function body " "inside a ```python code block. Do not include the function signature." ) # LiveCodeBench: AA uses original prompts without custom system prompts _LCB_SYSTEM = ( "You are an expert Python programmer. You will be given a question " "(problem specification) and will generate a correct Python program " "that matches the specification and passes all tests." ) _LCB_WITH_STARTER = ( "### Question:\n{question}\n\n" "### Format: You will use the following starter code to write the " "solution to the problem and enclose your code within delimiters.\n" "```python\n{starter_code}\n```\n\n" "### Answer: (use the provided format with backticks)\n" ) _LCB_WITHOUT_STARTER = ( "### Question:\n{question}\n\n" "### Format: Read the inputs from stdin solve the problem and write " "the answer to stdout (do not directly test on the sample inputs). " "Enclose your code within delimiters as follows. Ensure that when the " "python program runs, it reads the inputs, runs the algorithm and " "writes output to STDOUT.\n" "```python\n# YOUR CODE HERE\n```\n\n" "### Answer: (use the provided format with backticks)\n" ) def format_gpqa_question(doc: dict, idx: int) -> tuple[str, str]: """Returns (prompt, correct_letter).""" correct = doc["Correct Answer"] choices = [ correct, doc["Incorrect Answer 1"], doc["Incorrect Answer 2"], doc["Incorrect Answer 3"], ] rng = random.Random(idx) order = rng.sample(range(4), 4) shuffled = [choices[i] for i in order] correct_letter = "ABCD"[order.index(0)] choices_text = "\n".join(f"{L}) {shuffled[i]}" for i, L in enumerate("ABCD")) return f"{_GPQA_INSTRUCTION}\n\n{doc['Question']}\n\n{choices_text}", correct_letter def format_mmlu_pro_question(doc: dict) -> tuple[str, str]: """Returns (prompt, correct_letter).""" options = doc["options"] letters = "ABCDEFGHIJ" choices_text = "\n".join(f"{letters[i]}) {opt}" for i, opt in enumerate(options)) return f"{_MMLU_PRO_INSTRUCTION}\n\n{doc['question']}\n\n{choices_text}", doc[ "answer" ] def format_aime_question(doc: dict) -> tuple[str, int]: """Returns (prompt, correct_answer_int).""" return f"{_AIME_INSTRUCTION}\n\n{doc['problem']}", int(doc["answer"]) def format_humaneval_question(doc: dict) -> tuple[str, dict]: """Returns (prompt, metadata_for_execution).""" prompt = f"{_HUMANEVAL_INSTRUCTION}\n\n```python\n{doc['prompt']}```" # Pass the full problem dict — check_correctness needs task_id, prompt, # test, entry_point meta = { "problem": { "task_id": doc["task_id"], "prompt": doc["prompt"], "test": doc["test"], "entry_point": doc["entry_point"], }, } return prompt, meta def format_livecodebench_question(doc: dict) -> tuple[str, str | None, dict]: """Returns (prompt, system_message, metadata_for_execution).""" starter_code = doc.get("starter_code", "") question_content = doc["question_content"] if starter_code and starter_code.strip(): user_msg = _LCB_WITH_STARTER.format( question=question_content, starter_code=starter_code ) else: user_msg = _LCB_WITHOUT_STARTER.format(question=question_content) # Parse test cases public_tests = ( json.loads(doc["public_test_cases"]) if isinstance(doc["public_test_cases"], str) else doc["public_test_cases"] ) private_tests = doc.get("private_test_cases", "[]") if isinstance(private_tests, str): try: private_tests = json.loads(private_tests) except Exception: import base64 import pickle import zlib private_tests = json.loads( pickle.loads( zlib.decompress(base64.b64decode(private_tests.encode("utf-8"))) ) ) all_tests = public_tests + ( private_tests if isinstance(private_tests, list) else [] ) test_inputs = [t["input"] for t in all_tests] test_outputs = [t["output"] for t in all_tests] metadata = doc.get("metadata", "{}") if isinstance(metadata, str): metadata = json.loads(metadata) func_name = metadata.get("func_name") # Build the sample dict in official LCB format for run_test() input_output: dict[str, Any] = { "inputs": test_inputs, "outputs": test_outputs, } if func_name: input_output["fn_name"] = func_name meta = { "sample": {"input_output": json.dumps(input_output)}, } return user_msg, _LCB_SYSTEM, meta # --------------------------------------------------------------------------- # API client with retries # --------------------------------------------------------------------------- @dataclass class ApiResult: content: str prompt_tokens: int completion_tokens: int reasoning_tokens: int async def _call_api( client: httpx.AsyncClient, base_url: str, model: str, prompt: str, temperature: float, max_tokens: int, timeout: float | None, system_message: str | None = None, reasoning_effort: str | None = None, top_p: float | None = None, ) -> ApiResult: messages = [] if system_message: messages.append({"role": "system", "content": system_message}) messages.append({"role": "user", "content": prompt}) body: dict[str, Any] = { "model": model, "messages": messages, "temperature": temperature, "max_tokens": max_tokens, } if reasoning_effort is not None: body["reasoning_effort"] = reasoning_effort if top_p is not None: body["top_p"] = top_p resp = await client.post( f"{base_url}/v1/chat/completions", json=body, timeout=timeout, ) resp.raise_for_status() data = resp.json() content = data["choices"][0]["message"]["content"] if not content or not content.strip(): raise ValueError("Empty response from model") usage = data.get("usage", {}) details = usage.get("completion_tokens_details", {}) return ApiResult( content=content, prompt_tokens=usage.get("prompt_tokens", 0), completion_tokens=usage.get("completion_tokens", 0), reasoning_tokens=details.get("reasoning_tokens", 0) if details else 0, ) async def call_with_retries( client: httpx.AsyncClient, base_url: str, model: str, prompt: str, temperature: float, max_tokens: int, timeout: float | None = None, system_message: str | None = None, reasoning_effort: str | None = None, top_p: float | None = None, ) -> ApiResult | None: for attempt in range(MAX_RETRIES): try: return await _call_api( client, base_url, model, prompt, temperature, max_tokens, timeout, system_message, reasoning_effort, top_p, ) except Exception as e: if attempt < MAX_RETRIES - 1: wait = min(2**attempt, 60) logger.warning( f"Attempt {attempt + 1}/{MAX_RETRIES} failed: {e}. Retrying in {wait}s..." ) await asyncio.sleep(wait) else: logger.error(f"All {MAX_RETRIES} retries exhausted. Last error: {e}") return None # --------------------------------------------------------------------------- # Evaluation runners # --------------------------------------------------------------------------- async def evaluate_benchmark( benchmark_name: str, base_url: str, model: str, temperature: float, max_tokens: int, concurrency: int = 1, limit: int | None = None, timeout: float | None = None, reasoning_effort: str | None = None, top_p: float | None = None, difficulty: str | None = None, ) -> list[QuestionResult]: """Run a benchmark. Returns per-question results.""" import datasets config = BENCHMARKS[benchmark_name] logger.info(f"Loading dataset {config.dataset_name}...") try: if benchmark_name == "livecodebench": ds = datasets.load_dataset( "json", data_files="hf://datasets/livecodebench/code_generation_lite/*.jsonl", split="train", ) else: ds = datasets.load_dataset( config.dataset_name, config.dataset_config, split=config.split, ) except Exception as e: logger.error(f"Failed to load dataset: {e}") if "gated" in str(e).lower() or "login" in str(e).lower(): logger.error("Dataset requires authentication. Run: huggingface-cli login") return [] if difficulty and "difficulty" in ds.column_names: ds = ds.filter(lambda x: x["difficulty"] == difficulty) logger.info(f"Filtered to {len(ds)} {difficulty} problems") total = len(ds) if limit and limit < total: ds = ds.select(range(limit)) total = limit logger.info( f"Evaluating {benchmark_name}: {total} questions, concurrency={concurrency}, " f"temperature={temperature}, max_tokens={max_tokens}" ) if config.kind == "code": logger.warning( "Code benchmarks execute model-generated code. Use a sandboxed environment." ) semaphore = asyncio.Semaphore(concurrency) results: list[QuestionResult | None] = [None] * total completed = 0 lock = asyncio.Lock() async def process_question( idx: int, doc: dict, http_client: httpx.AsyncClient ) -> None: nonlocal completed system_msg = None if benchmark_name == "gpqa_diamond": prompt, gold = format_gpqa_question(doc, idx) valid_letters = "ABCD" elif benchmark_name == "mmlu_pro": prompt, gold = format_mmlu_pro_question(doc) valid_letters = "ABCDEFGHIJ"[: len(doc["options"])] elif benchmark_name.startswith("aime"): prompt, gold_int = format_aime_question(doc) gold = str(gold_int) elif benchmark_name == "humaneval": prompt, exec_meta = format_humaneval_question(doc) gold = "pass" elif benchmark_name == "livecodebench": prompt, system_msg, exec_meta = format_livecodebench_question(doc) gold = "pass" else: raise ValueError(f"Unknown benchmark: {benchmark_name}") async with semaphore: t0 = time.monotonic() api_result = await call_with_retries( http_client, base_url, model, prompt, temperature, max_tokens, timeout, system_message=system_msg, reasoning_effort=reasoning_effort, top_p=top_p, ) elapsed = time.monotonic() - t0 if api_result is None: result = QuestionResult( question_id=idx, prompt=prompt, response="", extracted_answer=None, gold_answer=gold, correct=False, error="API failure after retries", elapsed_s=elapsed, ) else: response = api_result.content stats = { "prompt_tokens": api_result.prompt_tokens, "completion_tokens": api_result.completion_tokens, "reasoning_tokens": api_result.reasoning_tokens, "elapsed_s": elapsed, } if config.kind == "mc": extracted = extract_mc_answer(response, valid_letters) result = QuestionResult( question_id=idx, prompt=prompt, response=response, extracted_answer=extracted, gold_answer=gold, correct=(extracted == gold) if extracted else False, **stats, ) elif config.kind == "math": extracted = extract_boxed_answer(response) correct = ( check_aime_answer(extracted, int(gold)) if extracted else False ) result = QuestionResult( question_id=idx, prompt=prompt, response=response, extracted_answer=extracted, gold_answer=gold, correct=correct, **stats, ) elif config.kind == "code": # HumanEval needs preserved indentation (function body completion) keep_indent = benchmark_name == "humaneval" code = extract_code_block(response, preserve_indent=keep_indent) if code is None: result = QuestionResult( question_id=idx, prompt=prompt, response=response, extracted_answer=None, gold_answer=gold, correct=False, error="No code block extracted", **stats, ) elif benchmark_name == "humaneval": passed, diag = run_humaneval_test( exec_meta["problem"], code, ) result = QuestionResult( question_id=idx, prompt=prompt, response=response, extracted_answer="pass" if passed else "fail", gold_answer=gold, correct=passed, error=diag if not passed else None, **stats, ) elif benchmark_name == "livecodebench": passed, diag = run_livecodebench_test( code, exec_meta["sample"], ) result = QuestionResult( question_id=idx, prompt=prompt, response=response, extracted_answer="pass" if passed else "fail", gold_answer=gold, correct=passed, error=diag if not passed else None, **stats, ) else: result = QuestionResult( question_id=idx, prompt=prompt, response=response, extracted_answer=None, gold_answer=gold, correct=False, error="Unknown code benchmark", **stats, ) else: result = QuestionResult( question_id=idx, prompt=prompt, response=response, extracted_answer=None, gold_answer=gold, correct=False, error="Unsupported kind", **stats, ) results[idx] = result async with lock: completed += 1 n = completed if n % max(1, total // 20) == 0 or n == total: correct_so_far = sum(1 for r in results if r is not None and r.correct) answered = sum(1 for r in results if r is not None) logger.info( f" [{n}/{total}] {correct_so_far}/{answered} correct " f"({correct_so_far / max(answered, 1):.1%})" ) async with httpx.AsyncClient() as http_client: tasks = [process_question(i, doc, http_client) for i, doc in enumerate(ds)] await asyncio.gather(*tasks) return [r for r in results if r is not None] # --------------------------------------------------------------------------- # Results display # --------------------------------------------------------------------------- def print_results( benchmark_name: str, results: list[QuestionResult], concurrency: int | None = None, ) -> dict[str, Any]: total = len(results) correct = sum(r.correct for r in results) errors = sum(1 for r in results if r.error) no_extract = sum(1 for r in results if r.extracted_answer is None and not r.error) accuracy = correct / max(total, 1) total_prompt_tokens = sum(r.prompt_tokens for r in results) total_completion_tokens = sum(r.completion_tokens for r in results) total_reasoning_tokens = sum(r.reasoning_tokens for r in results) total_elapsed = sum(r.elapsed_s for r in results) wall_clock = max(r.elapsed_s for r in results) if results else 0.0 avg_gen_tps = total_completion_tokens / total_elapsed if total_elapsed > 0 else 0.0 label = f"[c={concurrency}] " if concurrency is not None else "" print(f"\n{label}{benchmark_name}: {correct}/{total} ({accuracy:.1%})") tok_line = f" tokens: {total_prompt_tokens:,} prompt + {total_completion_tokens:,} completion" if total_reasoning_tokens > 0: tok_line += f" ({total_reasoning_tokens:,} reasoning)" tok_line += ( f" | avg gen tps: {avg_gen_tps:.1f}" f" | total time: {total_elapsed:.1f}s wall clock: {wall_clock:.1f}s" ) print(tok_line) if errors: print(f" API errors: {errors}") if no_extract: print(f" No answer extracted: {no_extract}") return { "benchmark": benchmark_name, "accuracy": accuracy, "correct": correct, "total": total, "errors": errors, "no_extract": no_extract, "total_prompt_tokens": total_prompt_tokens, "total_completion_tokens": total_completion_tokens, "total_reasoning_tokens": total_reasoning_tokens, "total_elapsed_s": total_elapsed, "wall_clock_s": wall_clock, "avg_gen_tps": avg_gen_tps, } def print_comparison( benchmark_name: str, results_by_c: dict[int, list[QuestionResult]], ) -> None: levels = sorted(results_by_c.keys()) print(f"\n{'=' * 70}") print(f"COMPARISON: {benchmark_name}") print(f"{'=' * 70}") header = f"{'Concurrency':<15} {'Accuracy':>10} {'Correct':>10} {'Total':>10} {'Comp Tokens':>12} {'Wall Clock':>12} {'Avg Gen TPS':>12}" print(header) print("-" * len(header)) for c in levels: r = results_by_c[c] correct = sum(q.correct for q in r) total = len(r) comp_tok = sum(q.completion_tokens for q in r) total_elapsed = sum(q.elapsed_s for q in r) avg_tps = comp_tok / total_elapsed if total_elapsed > 0 else 0.0 wall = max(q.elapsed_s for q in r) if r else 0.0 print( f"c={c:<13} {correct / max(total, 1):>10.1%} {correct:>10} {total:>10}" f" {comp_tok:>12,} {wall:>11.1f}s {avg_tps:>12.1f}" ) if len(levels) >= 2: base_results = results_by_c[levels[0]] test_results = results_by_c[levels[-1]] changed = sum( 1 for br, tr in zip(base_results, test_results, strict=True) if br.correct != tr.correct ) total = min(len(base_results), len(test_results)) print( f"\nQuestions with different correctness (c={levels[0]} vs c={levels[-1]}): {changed}/{total}" ) if changed == 0: print("Batching produced identical quality.") elif changed <= total * 0.01: print("Negligible quality difference from batching.") else: print( f"WARNING: {changed / max(total, 1) * 100:.1f}% of questions changed." ) print() # --------------------------------------------------------------------------- # Interactive task picker # --------------------------------------------------------------------------- def pick_tasks_interactive() -> list[str]: import termios import tty if not sys.stdin.isatty(): logger.error("No --tasks specified and stdin is not a terminal.") return [] items = [(name, cfg.description) for name, cfg in BENCHMARKS.items()] selected: list[bool] = [False] * len(items) cursor = 0 total_lines = len(items) + 4 def write(s: str) -> None: sys.stdout.write(s) def render(first: bool = False) -> None: if not first: write(f"\033[{total_lines}A") write("\033[J") write( "\033[1mSelect benchmarks\033[0m (up/down, space toggle, enter confirm, q quit)\r\n\r\n" ) for i, (name, desc) in enumerate(items): marker = ">" if i == cursor else " " check = "x" if selected[i] else " " line = f" {marker} [{check}] {name:<17} {desc}" write(f"\033[7m{line}\033[0m\r\n" if i == cursor else f"{line}\r\n") write(f"\r\n {sum(selected)} selected\r\n") sys.stdout.flush() fd = sys.stdin.fileno() old = termios.tcgetattr(fd) try: tty.setraw(fd) write("\033[?25l") render(first=True) while True: ch = sys.stdin.read(1) if ch in ("q", "\x03"): write("\033[?25h\033[0m\r\n") return [] elif ch in ("\r", "\n"): break elif ch == " ": selected[cursor] = not selected[cursor] elif ch == "\x1b": seq = sys.stdin.read(2) if seq == "[A": cursor = (cursor - 1) % len(items) elif seq == "[B": cursor = (cursor + 1) % len(items) render() finally: termios.tcsetattr(fd, termios.TCSADRAIN, old) write("\033[?25h\033[0m\r\n") sys.stdout.flush() chosen = [name for (name, _), sel in zip(items, selected, strict=True) if sel] if chosen: logger.info(f"Selected: {', '.join(chosen)}") return chosen # --------------------------------------------------------------------------- # Results persistence # --------------------------------------------------------------------------- def save_results( results_dir: str, benchmark_name: str, model: str, concurrency: int, results: list[QuestionResult], scores: dict[str, Any], ) -> Path: out_dir = Path(results_dir) / model.replace("/", "_") / benchmark_name out_dir.mkdir(parents=True, exist_ok=True) ts = time.strftime("%Y%m%d_%H%M%S") path = out_dir / f"c{concurrency}_{ts}.json" data = { "benchmark": benchmark_name, "model": model, "concurrency": concurrency, "scores": scores, "results": [ { "question_id": r.question_id, "prompt": r.prompt, "response": r.response, "extracted_answer": r.extracted_answer, "gold_answer": r.gold_answer, "correct": r.correct, "error": r.error, "prompt_tokens": r.prompt_tokens, "completion_tokens": r.completion_tokens, "reasoning_tokens": r.reasoning_tokens, "elapsed_s": round(r.elapsed_s, 2), } for r in results ], } with open(path, "w") as f: json.dump(data, f, indent=2) logger.info(f"Results saved to {path}") return path # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def parse_int_list(values: list[str]) -> list[int]: items: list[int] = [] for v in values: for part in v.split(","): if part.strip(): items.append(int(part.strip())) return items def main() -> int: ap = argparse.ArgumentParser( prog="exo-eval", description="Quality evaluation for exo — matches Artificial Analysis methodology.", ) add_common_instance_args(ap) ap.add_argument( "--tasks", default=None, help="Comma-separated benchmark names. Omit for interactive picker.", ) ap.add_argument( "--limit", type=int, default=None, help="Max questions per benchmark (for fast iteration).", ) reasoning_group = ap.add_mutually_exclusive_group() reasoning_group.add_argument( "--reasoning", action="store_true", default=None, help="Force reasoning-model settings (temperature=0.6, max_tokens=65536).", ) reasoning_group.add_argument( "--no-reasoning", action="store_true", default=False, help="Force non-reasoning settings (temperature=0, max_tokens=16384).", ) ap.add_argument( "--temperature", type=float, default=None, help="Override temperature." ) ap.add_argument("--top-p", type=float, default=None, help="Override top_p.") ap.add_argument( "--max-tokens", type=int, default=None, help="Override max output tokens." ) ap.add_argument( "--num-concurrent", type=int, default=1, help="Concurrent API requests (default: 1).", ) ap.add_argument( "--compare-concurrency", nargs="+", default=None, help="Run at multiple concurrency levels and compare. E.g. --compare-concurrency 1,4", ) ap.add_argument( "--request-timeout", type=float, default=None, help="Per-request timeout in seconds (default: no timeout).", ) ap.add_argument( "--reasoning-effort", default=None, choices=["low", "medium", "high"], help="Override reasoning effort (default: 'high' for reasoning models, none for non-reasoning).", ) ap.add_argument( "--difficulty", default=None, choices=["easy", "medium", "hard"], help="Filter by difficulty (livecodebench only). E.g. --difficulty hard", ) ap.add_argument( "--results-dir", default="eval_results", help="Directory for result JSON files (default: eval_results).", ) ap.add_argument( "--skip-instance-setup", action="store_true", help="Skip exo instance management (assumes model is already running).", ) args, _ = ap.parse_known_args() # Resolve tasks if args.tasks: task_names = [t.strip() for t in args.tasks.split(",") if t.strip()] else: task_names = pick_tasks_interactive() if not task_names: return 0 for t in task_names: if t not in BENCHMARKS: logger.error(f"Unknown benchmark '{t}'. Available: {', '.join(BENCHMARKS)}") return 1 # Instance management client = ExoClient(args.host, args.port, timeout_s=args.timeout) instance_id: str | None = None if not args.skip_instance_setup: short_id, full_model_id = resolve_model_short_id( client, args.model, force_download=args.force_download, ) selected = settle_and_fetch_placements( client, full_model_id, args, settle_timeout=args.settle_timeout, ) if not selected: logger.error("No valid placements matched your filters.") return 1 selected.sort( key=lambda p: ( str(p.get("instance_meta", "")), str(p.get("sharding", "")), -nodes_used_in_instance(p["instance"]), ), reverse=True, ) preview = selected[0] instance = preview["instance"] instance_id = instance_id_from_instance(instance) logger.info( f"PLACEMENT: {preview['sharding']} / {preview['instance_meta']} / " f"nodes={nodes_used_in_instance(instance)}" ) settle_deadline = ( time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None ) download_duration = run_planning_phase( client, full_model_id, preview, args.danger_delete_downloads, args.timeout, settle_deadline, ) if download_duration is not None: logger.info(f"Download: {download_duration:.1f}s") client.request_json("POST", "/instance", body={"instance": instance}) try: wait_for_instance_ready(client, instance_id) except (RuntimeError, TimeoutError) as e: logger.error(f"Failed to initialize: {e}") with contextlib.suppress(ExoHttpError): client.request_json("DELETE", f"/instance/{instance_id}") return 1 time.sleep(1) else: full_model_id = args.model # Auto-detect reasoning from model config model_config = load_model_config(full_model_id) if args.reasoning: is_reasoning = True elif args.no_reasoning: is_reasoning = False elif model_config is not None: is_reasoning = model_config.get("reasoning", False) logger.info( f"Auto-detected from config: {model_config['name']} → " f"{'reasoning' if is_reasoning else 'non-reasoning'}" ) else: is_reasoning = False logger.warning( f"Model '{full_model_id}' not found in eval_configs/models.toml. " f"Defaulting to non-reasoning. Use --reasoning to override." ) # Resolve temperature, max_tokens, reasoning_effort # Priority: CLI flag > per-model config > global defaults cfg = model_config or {} if args.temperature is not None: temperature = args.temperature elif "temperature" in cfg: temperature = float(cfg["temperature"]) else: temperature = ( TEMPERATURE_REASONING if is_reasoning else TEMPERATURE_NON_REASONING ) if args.max_tokens is not None: max_tokens = args.max_tokens elif "max_tokens" in cfg: max_tokens = int(cfg["max_tokens"]) else: max_tokens = REASONING_MAX_TOKENS if is_reasoning else DEFAULT_MAX_TOKENS if args.top_p is not None: top_p: float | None = args.top_p elif "top_p" in cfg: top_p = float(cfg["top_p"]) else: top_p = None # let server use its default if args.reasoning_effort is not None: reasoning_effort = args.reasoning_effort elif "reasoning_effort" in cfg: reasoning_effort = str(cfg["reasoning_effort"]) else: reasoning_effort = "high" if is_reasoning else None base_url = f"http://{args.host}:{args.port}" logger.info(f"Model: {full_model_id}") logger.info( f"Settings: temperature={temperature}, max_tokens={max_tokens}, " + (f"top_p={top_p}, " if top_p is not None else "") + f"reasoning={'yes' if is_reasoning else 'no'}" + (f", reasoning_effort={reasoning_effort}" if reasoning_effort else "") ) try: if args.compare_concurrency: concurrency_levels = parse_int_list(args.compare_concurrency) for task_name in task_names: results_by_c: dict[int, list[QuestionResult]] = {} for c in concurrency_levels: logger.info(f"\n{'=' * 50}") logger.info(f"Running {task_name} at concurrency={c}") results = asyncio.run( evaluate_benchmark( task_name, base_url, full_model_id, temperature, max_tokens, concurrency=c, limit=args.limit, timeout=args.request_timeout, reasoning_effort=reasoning_effort, top_p=top_p, difficulty=args.difficulty, ) ) if results: scores = print_results(task_name, results, concurrency=c) save_results( args.results_dir, task_name, full_model_id, c, results, scores, ) results_by_c[c] = results if len(results_by_c) >= 2: print_comparison(task_name, results_by_c) else: for task_name in task_names: results = asyncio.run( evaluate_benchmark( task_name, base_url, full_model_id, temperature, max_tokens, concurrency=args.num_concurrent, limit=args.limit, timeout=args.request_timeout, reasoning_effort=reasoning_effort, top_p=top_p, difficulty=args.difficulty, ) ) if results: scores = print_results(task_name, results) save_results( args.results_dir, task_name, full_model_id, args.num_concurrent, results, scores, ) finally: if instance_id is not None: try: client.request_json("DELETE", f"/instance/{instance_id}") except ExoHttpError as e: if e.status != 404: raise wait_for_instance_gone(client, instance_id) return 0 if __name__ == "__main__": raise SystemExit(main()) ================================================ FILE: bench/harness.py ================================================ # type: ignore from __future__ import annotations import argparse import http.client import json import os import time from typing import Any from urllib.parse import urlencode from loguru import logger _SETTLE_INITIAL_BACKOFF_S = 1.0 _SETTLE_MAX_BACKOFF_S = 60.0 _SETTLE_BACKOFF_MULTIPLIER = 2.0 class ExoHttpError(RuntimeError): def __init__(self, status: int, reason: str, body_preview: str): super().__init__(f"HTTP {status} {reason}: {body_preview}") self.status = status class ExoClient: def __init__(self, host: str, port: int, timeout_s: float = 7200.0): self.host = host self.port = port self.timeout_s = timeout_s def request_json( self, method: str, path: str, params: dict[str, Any] | None = None, body: dict[str, Any] | None = None, headers: dict[str, str] | None = None, ) -> Any: if not path.startswith("/"): path = "/" + path if params: path = path + "?" + urlencode(params) conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s) try: payload: bytes | None = None hdrs: dict[str, str] = {"Accept": "application/json"} if body is not None: payload = json.dumps(body).encode("utf-8") hdrs["Content-Type"] = "application/json" if headers: hdrs.update(headers) conn.request(method.upper(), path, body=payload, headers=hdrs) resp = conn.getresponse() raw = resp.read() text = raw.decode("utf-8", errors="replace") if raw else "" if resp.status >= 400: raise ExoHttpError(resp.status, resp.reason, text[:300]) if not text: return None return json.loads(text) finally: conn.close() def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]: return self.request_json("POST", "/bench/chat/completions", body=payload) def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]: if len(instance) != 1: raise KeyError(f"Expected 1 key, got keys={list(instance.keys())}") tag = next(iter(instance)) inner = instance[tag] if not isinstance(inner, dict): raise TypeError(f"payload for {tag} must be dict, got {type(inner)}") return inner def instance_id_from_instance(instance: dict[str, Any]) -> str: inner = unwrap_instance(instance) return str(inner["instanceId"]) def nodes_used_in_instance(instance: dict[str, Any]) -> int: inner = unwrap_instance(instance) return len(inner["shardAssignments"]["nodeToRunner"]) def runner_ids_from_instance(instance: dict[str, Any]) -> list[str]: inner = unwrap_instance(instance) runner_to_shard = inner["shardAssignments"]["runnerToShard"] return list(runner_to_shard.keys()) def runner_ready(runner: dict[str, Any]) -> bool: return "RunnerReady" in runner def runner_failed(runner: dict[str, Any]) -> bool: return "RunnerFailed" in runner def get_runner_failed_message(runner: dict[str, Any]) -> str | None: if "RunnerFailed" in runner: return runner["RunnerFailed"].get("errorMessage") return None def wait_for_instance_ready( client: ExoClient, instance_id: str, timeout: float = 24000.0 ) -> None: start_time = time.time() instance_existed = False while time.time() - start_time < timeout: state = client.request_json("GET", "/state") instances = state.get("instances", {}) if instance_id not in instances: if instance_existed: # Instance was deleted after being created - likely due to runner failure raise RuntimeError( f"Instance {instance_id} was deleted (runner may have failed)" ) time.sleep(0.1) continue instance_existed = True instance = instances[instance_id] runner_ids = runner_ids_from_instance(instance) runners = state.get("runners", {}) # Check for failed runners first for rid in runner_ids: runner = runners.get(rid, {}) if runner_failed(runner): error_msg = get_runner_failed_message(runner) or "Unknown error" raise RuntimeError(f"Runner {rid} failed: {error_msg}") if all(runner_ready(runners.get(rid, {})) for rid in runner_ids): return time.sleep(0.1) raise TimeoutError(f"Instance {instance_id} did not become ready within {timeout=}") def wait_for_instance_gone( client: ExoClient, instance_id: str, timeout: float = 3.0 ) -> None: start_time = time.time() while time.time() - start_time < timeout: try: client.request_json("GET", f"/instance/{instance_id}") time.sleep(0.4) except ExoHttpError as e: if e.status == 404: return raise raise TimeoutError(f"Instance {instance_id} did not get deleted within {timeout=}") def resolve_model_short_id( client: ExoClient, model_arg: str, *, force_download: bool = False ) -> tuple[str, str]: models = client.request_json("GET", "/models") or {} data = models.get("data") or [] for m in data: if (m.get("name") or "").lower() == model_arg.lower(): short_id = str(m["name"]) full_id = str(m.get("hugging_face_id") or m["name"]) return short_id, full_id for m in data: if m.get("hugging_face_id") == model_arg: short_id = str(m["name"]) full_id = str(m["hugging_face_id"]) return short_id, full_id if force_download and "/" in model_arg: logger.info(f"Model not in /models, adding from HuggingFace: {model_arg}") result = client.request_json( "POST", "/models/add", body={"model_id": model_arg} ) if result: short_id = str(result.get("name") or model_arg.rsplit("/", 1)[-1]) full_id = str(result.get("hugging_face_id") or model_arg) return short_id, full_id raise ValueError(f"Model not found in /models: {model_arg}") def placement_filter(instance_meta: str, wanted: str) -> bool: s = (instance_meta or "").lower() if wanted == "both": return ("ring" in s) or ("jaccl" in s) return wanted in s def sharding_filter(sharding: str, wanted: str) -> bool: s = (sharding or "").lower() if wanted == "both": return ("pipeline" in s) or ("tensor" in s) return wanted in s def fetch_and_filter_placements( client: ExoClient, full_model_id: str, args: argparse.Namespace ) -> list[dict[str, Any]]: previews_resp = client.request_json( "GET", "/instance/previews", params={"model_id": full_model_id} ) previews = previews_resp.get("previews") or [] selected: list[dict[str, Any]] = [] for p in previews: if p.get("error") is not None: continue if not placement_filter(str(p.get("instance_meta", "")), args.instance_meta): continue if not sharding_filter(str(p.get("sharding", "")), args.sharding): continue instance = p.get("instance") if not isinstance(instance, dict): continue n = nodes_used_in_instance(instance) # Skip tensor ring single node as it is pointless when pipeline ring if n == 1 and ( (args.sharding == "both" and "tensor" in p.get("sharding", "").lower()) or ( args.instance_meta == "both" and "jaccl" in p.get("instance_meta", "").lower() ) ): continue if ( args.skip_pipeline_jaccl and ( args.instance_meta == "both" and "jaccl" in p.get("instance_meta", "").lower() ) and ( args.sharding == "both" and "pipeline" in p.get("sharding", "").lower() ) ): continue if ( args.skip_tensor_ring and ( args.instance_meta == "both" and "ring" in p.get("instance_meta", "").lower() ) and (args.sharding == "both" and "tensor" in p.get("sharding", "").lower()) ): continue if args.min_nodes <= n <= args.max_nodes: selected.append(p) return selected def settle_and_fetch_placements( client: ExoClient, full_model_id: str, args: argparse.Namespace, settle_timeout: float = 0, ) -> list[dict[str, Any]]: selected = fetch_and_filter_placements(client, full_model_id, args) if not selected and settle_timeout > 0: backoff = _SETTLE_INITIAL_BACKOFF_S deadline = time.monotonic() + settle_timeout while not selected and time.monotonic() < deadline: remaining = deadline - time.monotonic() logger.warning( f"No valid placements yet (cluster may still be settling). " f"Retrying in {backoff:.1f}s ({remaining:.0f}s remaining)..." ) time.sleep(min(backoff, remaining)) backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S) selected = fetch_and_filter_placements(client, full_model_id, args) return selected def run_planning_phase( client: ExoClient, full_model_id: str, preview: dict[str, Any], danger_delete: bool, timeout: float, settle_deadline: float | None, ) -> float | None: """Check disk space and ensure model is downloaded before benchmarking. Returns the wall-clock download duration in seconds if a fresh download was needed, or None if the model was already cached on all nodes. """ # Get model size from /models models = client.request_json("GET", "/models") or {} model_bytes = 0 for m in models.get("data", []): if m.get("hugging_face_id") == full_model_id: model_bytes = m.get("storage_size_megabytes", 0) * 1024 * 1024 break if not model_bytes: logger.warning( f"Could not determine size for {full_model_id}, skipping disk check" ) return None # Get nodes from preview inner = unwrap_instance(preview["instance"]) node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys()) runner_to_shard = inner["shardAssignments"]["runnerToShard"] state = client.request_json("GET", "/state") downloads = state.get("downloads", {}) node_disk = state.get("nodeDisk", {}) needs_download = False for node_id in node_ids: node_downloads = downloads.get(node_id, []) # Check if model already downloaded on this node already_downloaded = any( "DownloadCompleted" in p and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][ "modelId" ] == full_model_id for p in node_downloads ) if already_downloaded: continue needs_download = True # Wait for disk info if settle_deadline is set disk_info = node_disk.get(node_id, {}) backoff = _SETTLE_INITIAL_BACKOFF_S while not disk_info and settle_deadline and time.monotonic() < settle_deadline: remaining = settle_deadline - time.monotonic() logger.info( f"Waiting for disk info on {node_id} ({remaining:.0f}s remaining)..." ) time.sleep(min(backoff, remaining)) backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S) state = client.request_json("GET", "/state") node_disk = state.get("nodeDisk", {}) disk_info = node_disk.get(node_id, {}) if not disk_info: logger.warning(f"No disk info for {node_id}, skipping space check") continue avail = disk_info.get("available", {}).get("inBytes", 0) if avail >= model_bytes: continue if not danger_delete: raise RuntimeError( f"Insufficient disk on {node_id}: need {model_bytes // (1024**3)}GB, " f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space." ) # Delete from smallest to largest (skip read-only models from EXO_MODELS_PATH) completed = [ ( unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][ "modelId" ], p["DownloadCompleted"]["total"]["inBytes"], ) for p in node_downloads if "DownloadCompleted" in p and not p["DownloadCompleted"].get("readOnly", False) ] for del_model, size in sorted(completed, key=lambda x: x[1]): logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)") client.request_json("DELETE", f"/download/{node_id}/{del_model}") avail += size if avail >= model_bytes: break if avail < model_bytes: raise RuntimeError(f"Could not free enough space on {node_id}") # Start downloads (idempotent) download_t0 = time.perf_counter() if needs_download else None for node_id in node_ids: runner_id = inner["shardAssignments"]["nodeToRunner"][node_id] shard = runner_to_shard[runner_id] client.request_json( "POST", "/download/start", body={ "targetNodeId": node_id, "shardMetadata": shard, }, ) logger.info(f"Started download on {node_id}") # Wait for downloads start = time.time() while time.time() - start < timeout: state = client.request_json("GET", "/state") downloads = state.get("downloads", {}) all_done = True for node_id in node_ids: done = any( "DownloadCompleted" in p and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[ "modelCard" ]["modelId"] == full_model_id for p in downloads.get(node_id, []) ) failed = [ p["DownloadFailed"]["errorMessage"] for p in downloads.get(node_id, []) if "DownloadFailed" in p and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][ "modelId" ] == full_model_id ] if failed: raise RuntimeError(f"Download failed on {node_id}: {failed[0]}") if not done: all_done = False if all_done: if download_t0 is not None: return time.perf_counter() - download_t0 return None time.sleep(1) raise TimeoutError("Downloads did not complete in time") def add_common_instance_args(ap: argparse.ArgumentParser) -> None: ap.add_argument("--host", default=os.environ.get("EXO_HOST", "localhost")) ap.add_argument( "--port", type=int, default=int(os.environ.get("EXO_PORT", "52415")) ) ap.add_argument("--model", required=True, help="Model short id or huggingface id") ap.add_argument( "--force-download", action="store_true", help="If model not in /models, add it from HuggingFace via exo and download.", ) ap.add_argument( "--max-nodes", type=int, default=4, help="Only consider placements using <= this many nodes.", ) ap.add_argument( "--min-nodes", type=int, default=1, help="Only consider placements using >= this many nodes.", ) ap.add_argument( "--instance-meta", choices=["ring", "jaccl", "both"], default="both" ) ap.add_argument( "--sharding", choices=["pipeline", "tensor", "both"], default="both" ) ap.add_argument( "--skip-pipeline-jaccl", action="store_true", help="Skip pipeline+jaccl placements, as it's often pointless.", ) ap.add_argument( "--skip-tensor-ring", action="store_true", help="Skip tensor+ring placements, as it's so slow.", ) ap.add_argument( "--timeout", type=float, default=7200.0, help="HTTP timeout (seconds)." ) ap.add_argument( "--settle-timeout", type=float, default=0, help="Max seconds to wait for the cluster to produce valid placements (0 = try once).", ) ap.add_argument( "--danger-delete-downloads", action="store_true", help="Delete existing models from smallest to largest to make room for benchmark model.", ) ================================================ FILE: bench/parallel_requests.py ================================================ # type: ignore import argparse import asyncio import sys import termios import time import tty import aiohttp NUM_REQUESTS = 10 BASE_URL = "" QUESTIONS = [ "What is the capital of Australia?", "How many bones are in the human body?", "What year did World War II end?", "What is the speed of light in meters per second?", "Who wrote Romeo and Juliet?", "What is the chemical formula for water?", "How many planets are in our solar system?", "What is the largest ocean on Earth?", "Who painted the Mona Lisa?", "What is the boiling point of water in Celsius?", ] def write(s: str) -> None: sys.stdout.write(s) # --------------------------------------------------------------------------- # Model picker (same style as exo_eval) # --------------------------------------------------------------------------- def fetch_models() -> list[str]: import json import urllib.request with urllib.request.urlopen(f"{BASE_URL}/state") as resp: data = json.loads(resp.read()) model_ids: set[str] = set() for instance in data.get("instances", {}).values(): for variant in instance.values(): sa = variant.get("shardAssignments", {}) model_id = sa.get("modelId") if model_id: model_ids.add(model_id) return sorted(model_ids) def pick_model() -> str | None: models = fetch_models() if not models: print("No models found.") return None cursor = 0 total_lines = len(models) + 4 def render(first: bool = False) -> None: if not first: write(f"\033[{total_lines}A") write("\033[J") write("\033[1mSelect model\033[0m (up/down, enter confirm, q quit)\r\n\r\n") for i, model in enumerate(models): line = f" {'>' if i == cursor else ' '} {model}" write(f"\033[7m{line}\033[0m\r\n" if i == cursor else f"{line}\r\n") write("\r\n") sys.stdout.flush() fd = sys.stdin.fileno() old = termios.tcgetattr(fd) try: tty.setraw(fd) write("\033[?25l") render(first=True) while True: ch = sys.stdin.read(1) if ch in ("q", "\x03"): write("\033[?25h\033[0m\r\n") return None elif ch in ("\r", "\n"): break elif ch == "\x1b": seq = sys.stdin.read(2) if seq == "[A": cursor = (cursor - 1) % len(models) elif seq == "[B": cursor = (cursor + 1) % len(models) render() finally: termios.tcsetattr(fd, termios.TCSADRAIN, old) write(f"\033[{total_lines}A\033[J") # clear picker UI write("\033[?25h\033[0m") sys.stdout.flush() return models[cursor] # --------------------------------------------------------------------------- # Parallel requests # --------------------------------------------------------------------------- statuses: list[str] = [] times: list[str] = [] previews: list[str] = [] tokens: list[str] = [] full_responses: list[dict | None] = [] total_lines = 0 start_time: float = 0 selected_model: str = "" def render_progress(first: bool = False) -> None: if not first: write(f"\033[{total_lines}A") write("\033[J") elapsed = time.monotonic() - start_time if start_time else 0 done = sum(1 for s in statuses if s == "done") write( f"\033[1m{selected_model}\033[0m [{done}/{NUM_REQUESTS}] {elapsed:.1f}s\r\n\r\n" ) for i in range(NUM_REQUESTS): q = QUESTIONS[i % len(QUESTIONS)] status = statuses[i] if status == "pending": color = "\033[33m" # yellow elif status == "running": color = "\033[36m" # cyan elif status == "done": color = "\033[32m" # green else: color = "\033[31m" # red write( f" {i:>2} {color}{status:<8}\033[0m {times[i]:>6} {tokens[i]:>5}tok {q[:40]:<40} {previews[i][:50]}\r\n" ) write("\r\n") sys.stdout.flush() async def send_request( session: aiohttp.ClientSession, i: int, lock: asyncio.Lock ) -> None: payload = { "model": selected_model, "messages": [{"role": "user", "content": QUESTIONS[i % len(QUESTIONS)]}], "max_tokens": 1024, } statuses[i] = "running" async with lock: render_progress() t0 = time.monotonic() try: async with session.post( f"{BASE_URL}/v1/chat/completions", json=payload ) as resp: data = await resp.json() elapsed = time.monotonic() - t0 full_responses[i] = data times[i] = f"{elapsed:.1f}s" if resp.status == 200: choice = data["choices"][0] msg = choice["message"] content = msg.get("content", "") previews[i] = content[:50].replace("\n", " ") or "(empty)" if "usage" in data: tokens[i] = str(data["usage"].get("total_tokens", "")) statuses[i] = "done" else: statuses[i] = f"err:{resp.status}" previews[i] = str(data.get("error", {}).get("message", ""))[:50] except Exception as e: elapsed = time.monotonic() - t0 times[i] = f"{elapsed:.1f}s" statuses[i] = "error" previews[i] = str(e)[:50] async with lock: render_progress() async def run_requests(print_stdout: bool = False) -> None: global start_time, total_lines, statuses, times, previews, tokens, full_responses statuses = ["pending"] * NUM_REQUESTS times = ["-"] * NUM_REQUESTS previews = ["-"] * NUM_REQUESTS tokens = ["-"] * NUM_REQUESTS full_responses = [None] * NUM_REQUESTS total_lines = NUM_REQUESTS + 4 write("\033[?25l") # hide cursor start_time = time.monotonic() render_progress(first=True) lock = asyncio.Lock() try: async with aiohttp.ClientSession() as session: tasks = [send_request(session, i, lock) for i in range(NUM_REQUESTS)] await asyncio.gather(*tasks) total = time.monotonic() - start_time write( f"\033[1m=== All {NUM_REQUESTS} requests done in {total:.1f}s ===\033[0m\r\n\r\n" ) if print_stdout: for i in range(NUM_REQUESTS): data = full_responses[i] if not data or "choices" not in data: continue choice = data["choices"][0] msg = choice["message"] q = QUESTIONS[i % len(QUESTIONS)] write(f"\033[1m--- #{i}: {q} ---\033[0m\r\n") if msg.get("reasoning_content"): write(f"\033[2m[Thinking]: {msg['reasoning_content']}\033[0m\r\n") write(f"{msg.get('content', '')}\r\n") if "usage" in data: u = data["usage"] write( f"\033[2m[Usage: prompt={u.get('prompt_tokens')}, " f"completion={u.get('completion_tokens')}, " f"total={u.get('total_tokens')}]\033[0m\r\n" ) write("\r\n") finally: write("\033[?25h") # show cursor sys.stdout.flush() def main() -> None: global selected_model, BASE_URL parser = argparse.ArgumentParser( description="Send parallel requests to an exo cluster" ) parser.add_argument( "--host", required=True, help="Hostname of the exo node (e.g. s1)" ) parser.add_argument("--port", type=int, default=52415, help="Port (default: 52415)") parser.add_argument( "--stdout", action="store_true", help="Print full responses after completion" ) args = parser.parse_args() BASE_URL = f"http://{args.host}:{args.port}" model = pick_model() if not model: return selected_model = model asyncio.run(run_requests(print_stdout=args.stdout)) if __name__ == "__main__": main() ================================================ FILE: bench/pyproject.toml ================================================ [project] name = "exo-bench" version = "0.1.0" description = "Benchmarking tool for exo distributed inference" requires-python = ">=3.13" dependencies = [ "httpx>=0.27.0", "loguru>=0.7.3", "transformers>=5.0.0", "huggingface-hub>=0.33.4", "tiktoken>=0.12.0", "jinja2>=3.1.0", "protobuf>=5.29.0", "datasets>=2.0.0", "math-verify>=0.7.0", "lm-eval[api,math]>=0.4.0", "human-eval>=1.0.3", "numpy>=1.24.0", ] [build-system] requires = ["hatchling"] build-backend = "hatchling.build" ================================================ FILE: bench/scenarios.toml ================================================ # Tool definitions — each becomes an OpenAI function tool. # All scenarios get all tools unless they specify a `tools` list. [tools.get_current_weather] description = "Get the current weather in a given location" required = ["location"] [tools.get_current_weather.properties.location] type = "string" description = "City and state, e.g. San Francisco, CA" [tools.get_current_weather.properties.unit] type = "string" enum = ["celsius", "fahrenheit"] description = "Temperature unit" [tools.calculate] description = "Evaluate a mathematical expression and return the numeric result" required = ["expression"] [tools.calculate.properties.expression] type = "string" description = "The math expression to evaluate, e.g. '2 + 3 * 4'" [tools.search_products] description = "Search for products in a catalog by query, category, and price" required = ["query"] [tools.search_products.properties.query] type = "string" description = "Search query string" [tools.search_products.properties.category] type = "string" enum = ["electronics", "clothing", "food", "books"] description = "Product category to filter by" [tools.search_products.properties.max_price] type = "number" description = "Maximum price in USD" [tools.create_todos] description = "Create a structured todo list" required = ["todos"] [tools.create_todos.properties.todos] type = "array" description = "List of todo items" [tools.create_todos.properties.todos.items] type = "object" required = ["content", "status", "priority"] [tools.create_todos.properties.todos.items.properties.content] type = "string" description = "The todo item text" [tools.create_todos.properties.todos.items.properties.status] type = "string" description = "Status: pending, in_progress, or completed" [tools.create_todos.properties.todos.items.properties.priority] type = "string" description = "Priority: low, normal, or high" # -- Should call a tool -- [[scenarios]] name = "weather_simple" description = "Basic weather query -> get_current_weather" expect_tool_call = true expected_function = "get_current_weather" required_arg_keys = ["location"] [[scenarios.messages]] role = "user" content = "What's the weather like in Tokyo right now?" [[scenarios]] name = "calculator_simple" description = "Math question -> calculate" expect_tool_call = true expected_function = "calculate" required_arg_keys = ["expression"] [[scenarios.messages]] role = "user" content = "Use the calculator to compute 3847 * 926 + 17293" [[scenarios]] name = "search_with_filters" description = "Product search with category and price filter" expect_tool_call = true expected_function = "search_products" required_arg_keys = ["query"] [[scenarios.messages]] role = "user" content = "Find me electronics under $50" # -- Multi-turn: tool call then follow-up -- [[scenarios]] name = "weather_multi_turn" description = "Weather query -> tool result -> natural language summary" expect_tool_call = true expected_function = "get_current_weather" required_arg_keys = ["location"] [scenarios.tool_result] temperature = "18C" condition = "partly cloudy" humidity = "65%" wind = "12 km/h NW" [[scenarios.messages]] role = "user" content = "What's the weather in Paris?" [[scenarios]] name = "calculator_multi_turn" description = "Math query -> tool result -> model reports the answer" expect_tool_call = true expected_function = "calculate" required_arg_keys = ["expression"] [scenarios.tool_result] result = 491682 [[scenarios.messages]] role = "user" content = "Use the calculator to compute 1847 * 263 + 5921" [[scenarios]] name = "search_multi_turn" description = "Search query -> tool result -> model summarizes products" expect_tool_call = true expected_function = "search_products" required_arg_keys = ["query"] [[scenarios.tool_result.results]] name = "Hands-On Machine Learning" price = 45.99 rating = 4.8 [[scenarios.tool_result.results]] name = "Deep Learning with Python" price = 39.99 rating = 4.6 [[scenarios.messages]] role = "user" content = "Search for books about machine learning" # -- Sequential tool calls -- [[scenarios]] name = "chained_tool_calls_same" description = "Thinking + weather(Tokyo) -> result -> model must call weather(London)" expect_tool_call = true expected_function = "get_current_weather" required_arg_keys = ["location"] [[scenarios.messages]] role = "user" content = "Compare the weather in Tokyo and London." [[scenarios.messages]] role = "assistant" content = "I'll check both cities. Let me start with Tokyo." [[scenarios.messages.tool_calls]] id = "call_1" name = "get_current_weather" arguments = { location = "Tokyo" } [[scenarios.messages]] role = "tool" tool_call_id = "call_1" content = '{"temperature": "25C", "condition": "sunny"}' [[scenarios]] name = "chained_tool_calls_different" description = "Thinking + weather(Berlin) -> result -> model must call calculator" expect_tool_call = true expected_function = "calculate" required_arg_keys = ["expression"] [[scenarios.messages]] role = "user" content = "What's the weather in Berlin, and also use the calculator to compute 4819 * 37 + 291." [[scenarios.messages]] role = "assistant" content = "I'll handle both. Let me check Berlin's weather first." [[scenarios.messages.tool_calls]] id = "call_2" name = "get_current_weather" arguments = { location = "Berlin" } [[scenarios.messages]] role = "tool" tool_call_id = "call_2" content = '{"temperature": "12C", "condition": "rainy"}' [[scenarios]] name = "chained_tool_calls_three" description = "Two prior thinking+tool calls -> results -> model must make a third" expect_tool_call = true expected_function = "get_current_weather" required_arg_keys = ["location"] [[scenarios.messages]] role = "user" content = "Compare weather in Tokyo, Paris, and London." [[scenarios.messages]] role = "assistant" content = "I'll check all three cities. Starting with Tokyo." [[scenarios.messages.tool_calls]] id = "call_3" name = "get_current_weather" arguments = { location = "Tokyo" } [[scenarios.messages]] role = "tool" tool_call_id = "call_3" content = '{"temperature": "25C", "condition": "sunny"}' [[scenarios.messages]] role = "assistant" content = "Got Tokyo. Now checking Paris." [[scenarios.messages.tool_calls]] id = "call_4" name = "get_current_weather" arguments = { location = "Paris" } [[scenarios.messages]] role = "tool" tool_call_id = "call_4" content = '{"temperature": "18C", "condition": "cloudy"}' # -- Nested object schema (regression for lossy chat template rendering) -- [[scenarios]] name = "nested_schema_tool_call" description = "Tool call with nested object array schema -> create_todos" expect_tool_call = true expected_function = "create_todos" required_arg_keys = ["todos"] nested_array_key = "todos" required_item_keys = ["content", "status", "priority"] tools = ["create_todos"] [[scenarios.messages]] role = "user" content = "Create a todo list with 3 items to learn Python" # -- Tool name integrity (regression for harmony token leaking into name) -- [tools.glob] description = "Search for files matching a glob pattern in the codebase" required = ["pattern"] [tools.glob.properties.pattern] type = "string" description = "The glob pattern to match files against, e.g. '**/*.py'" [tools.glob.properties.path] type = "string" description = "The directory to search in" [[scenarios]] name = "tool_name_integrity" description = "Tool name must not contain harmony tokens like <|channel|>" expect_tool_call = true expected_function = "glob" required_arg_keys = ["pattern"] tools = ["glob"] [[scenarios.messages]] role = "user" content = "Find all Python files in the src directory" # -- Should NOT call a tool -- [[scenarios]] name = "no_tool_joke" description = "Joke request should NOT trigger any tool" expect_tool_call = false [[scenarios.messages]] role = "user" content = "Tell me a funny joke about cats." [[scenarios]] name = "no_tool_factual" description = "Factual question answerable from training data" expect_tool_call = false [[scenarios.messages]] role = "user" content = "What is the capital of Japan?" ================================================ FILE: bench/single-m3-ultra.toml ================================================ # Single-node M3 Ultra benchmarks # # Shared constraints applied to ALL benchmarks in this file. constraints = [ "All(MacOsBuild(=25D125))", "Hosts(=1)", "All(Chip(m3_ultra))", "All(GpuCores(=80))", ] [topology] type = "none" # Default args merged into each benchmark's args (benchmark-level args win). [defaults] pp = [512, 2048, 8192, 16384] tg = 128 [[benchmark]] model = "mlx-community/Meta-Llama-3.1-70B-Instruct-4bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/gpt-oss-120b-MXFP4-Q8" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/GLM-4.7-Flash-8bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Qwen3-Coder-Next-6bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Qwen3-30B-A3B-8bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Qwen3-0.6B-4bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Qwen3-0.6B-8bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Llama-3.2-1B-Instruct-4bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Llama-3.2-3B-Instruct-4bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Llama-3.2-3B-Instruct-8bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Meta-Llama-3.1-8B-Instruct-8bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Meta-Llama-3.1-8B-Instruct-bf16" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/gpt-oss-20b-MXFP4-Q8" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Qwen3-30B-A3B-4bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/GLM-4.7-Flash-4bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/GLM-4.7-Flash-5bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/GLM-4.7-Flash-6bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Llama-3.3-70B-Instruct-4bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Qwen3-Coder-Next-4bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Qwen3-Coder-Next-5bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Qwen3-Coder-Next-8bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Qwen3-Next-80B-A3B-Instruct-4bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Qwen3-Next-80B-A3B-Instruct-8bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Qwen3-Next-80B-A3B-Thinking-4bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Qwen3-Next-80B-A3B-Thinking-8bit" extra_constraints = ["All(Memory(>=96GiB))"] [[benchmark]] model = "mlx-community/Llama-3.3-70B-Instruct-8bit" extra_constraints = ["All(Memory(>=256GiB))"] [[benchmark]] model = "mlx-community/llama-3.3-70b-instruct-fp16" extra_constraints = ["All(Memory(>=256GiB))"] [[benchmark]] model = "mlx-community/GLM-4.5-Air-8bit" extra_constraints = ["All(Memory(>=256GiB))"] [[benchmark]] model = "mlx-community/GLM-4.5-Air-bf16" extra_constraints = ["All(Memory(>=256GiB))"] [[benchmark]] model = "mlx-community/GLM-4.7-4bit" extra_constraints = ["All(Memory(>=256GiB))"] [[benchmark]] model = "mlx-community/MiniMax-M2.1-3bit" extra_constraints = ["All(Memory(>=256GiB))"] [[benchmark]] model = "mlx-community/MiniMax-M2.1-8bit" extra_constraints = ["All(Memory(>=256GiB))"] [[benchmark]] model = "mlx-community/Qwen3-235B-A22B-Instruct-2507-4bit" extra_constraints = ["All(Memory(>=256GiB))"] [[benchmark]] model = "mlx-community/Qwen3-Coder-Next-bf16" extra_constraints = ["All(Memory(>=256GiB))"] [[benchmark]] model = "mlx-community/Step-3.5-Flash-4bit" extra_constraints = ["All(Memory(>=256GiB))"] [[benchmark]] model = "mlx-community/Step-3.5-Flash-6bit" extra_constraints = ["All(Memory(>=256GiB))"] [[benchmark]] model = "mlx-community/Step-3.5-Flash-8Bit" extra_constraints = ["All(Memory(>=256GiB))"] [[benchmark]] model = "mlx-community/DeepSeek-V3.1-4bit" extra_constraints = ["All(Memory(>=512GiB))"] [[benchmark]] model = "mlx-community/GLM-4.7-6bit" extra_constraints = ["All(Memory(>=512GiB))"] [[benchmark]] model = "mlx-community/GLM-4.7-8bit-gs32" extra_constraints = ["All(Memory(>=512GiB))"] [[benchmark]] model = "mlx-community/Qwen3-235B-A22B-Instruct-2507-8bit" extra_constraints = ["All(Memory(>=512GiB))"] [[benchmark]] model = "mlx-community/Qwen3-Coder-480B-A35B-Instruct-4bit" extra_constraints = ["All(Memory(>=512GiB))"] ================================================ FILE: bench/src/exo_bench/__init__.py ================================================ ================================================ FILE: bench/vendor/__init__.py ================================================ ================================================ FILE: bench/vendor/lcb_testing_util.py ================================================ # type: ignore # Vendored from LiveCodeBench (https://github.com/LiveCodeBench/LiveCodeBench) # File: lcb_runner/evaluation/testing_util.py # License: MIT # Vendored 2026-03-07 — do not modify without updating from upstream. import ast import faulthandler import json import platform # to run the solution files we're using a timing based approach import signal import sys import time # used for debugging to time steps from datetime import datetime from decimal import Decimal from enum import Enum from io import StringIO # from pyext import RuntimeModule from types import ModuleType # used for testing the code that reads from input from unittest.mock import mock_open, patch import_string = "from string import *\nfrom re import *\nfrom datetime import *\nfrom collections import *\nfrom heapq import *\nfrom bisect import *\nfrom copy import *\nfrom math import *\nfrom random import *\nfrom statistics import *\nfrom itertools import *\nfrom functools import *\nfrom operator import *\nfrom io import *\nfrom sys import *\nfrom json import *\nfrom builtins import *\nfrom typing import *\nimport string\nimport re\nimport datetime\nimport collections\nimport heapq\nimport bisect\nimport copy\nimport math\nimport random\nimport statistics\nimport itertools\nimport functools\nimport operator\nimport io\nimport sys\nimport json\nsys.setrecursionlimit(50000)\n" def truncatefn(s, length=300): if isinstance(s, str): pass else: s = str(s) if len(s) <= length: return s return s[: length // 2] + "...(truncated) ..." + s[-length // 2 :] class CODE_TYPE(Enum): call_based = 0 standard_input = 1 # stuff for setting up signal timer class TimeoutException(Exception): pass def timeout_handler(signum, frame): print("timeout occured: alarm went off") raise TimeoutException # used to capture stdout as a list # from https://stackoverflow.com/a/16571630/6416660 # alternative use redirect_stdout() from contextlib class Capturing(list): def __enter__(self): self._stdout = sys.stdout sys.stdout = self._stringio = StringIO() # Make closing the StringIO a no-op self._stringio.close = lambda x: 1 return self def __exit__(self, *args): self.append(self._stringio.getvalue()) del self._stringio # free up some memory sys.stdout = self._stdout # Custom mock for sys.stdin that supports buffer attribute class MockStdinWithBuffer: def __init__(self, inputs: str): self.inputs = inputs self._stringio = StringIO(inputs) self.buffer = MockBuffer(inputs) def read(self, *args): return self.inputs def readline(self, *args): return self._stringio.readline(*args) def readlines(self, *args): return self.inputs.split("\n") def __getattr__(self, name): # Delegate other attributes to StringIO return getattr(self._stringio, name) class MockBuffer: def __init__(self, inputs: str): self.inputs = inputs.encode("utf-8") # Convert to bytes def read(self, *args): # Return as byte strings that can be split return self.inputs def readline(self, *args): return self.inputs.split(b"\n")[0] + b"\n" def clean_if_name(code: str) -> str: try: astree = ast.parse(code) last_block = astree.body[-1] if isinstance(last_block, ast.If): condition = last_block.test if ast.unparse(condition).strip() == "__name__ == '__main__'": code = ( ast.unparse(astree.body[:-1]) + "\n" + ast.unparse(last_block.body) # type: ignore ) except: pass return code def make_function(code: str) -> str: try: import_stmts = [] all_other_stmts = [] astree = ast.parse(code) for stmt in astree.body: if isinstance(stmt, (ast.Import, ast.ImportFrom)): import_stmts.append(stmt) else: all_other_stmts.append(stmt) function_ast = ast.FunctionDef( name="wrapped_function", args=ast.arguments( posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[] ), body=all_other_stmts, decorator_list=[], lineno=-1, ) main_code = ( import_string + "\n" + ast.unparse(import_stmts) + "\n" + ast.unparse(function_ast) ) return main_code except Exception: return code def call_method(method, inputs): if isinstance(inputs, list): inputs = "\n".join(inputs) inputs_line_iterator = iter(inputs.split("\n")) # Create custom stdin mock with buffer support mock_stdin = MockStdinWithBuffer(inputs) # sys.setrecursionlimit(10000) # @patch('builtins.input', side_effect=inputs.split("\n")) @patch("builtins.open", mock_open(read_data=inputs)) @patch("sys.stdin", mock_stdin) # Use our custom mock instead of StringIO @patch("sys.stdin.readline", lambda *args: next(inputs_line_iterator)) @patch("sys.stdin.readlines", lambda *args: inputs.split("\n")) @patch("sys.stdin.read", lambda *args: inputs) # @patch('sys.stdout.write', print) def _inner_call_method(_method): try: return _method() except SystemExit: pass finally: pass return _inner_call_method(method) def get_function(compiled_sol, fn_name: str): # type: ignore try: assert hasattr(compiled_sol, fn_name) return getattr(compiled_sol, fn_name) except Exception: return def compile_code(code: str, timeout: int): signal.alarm(timeout) try: tmp_sol = ModuleType("tmp_sol", "") exec(code, tmp_sol.__dict__) if "class Solution" in code: # leetcode wraps solutions in `Solution` # this is a hack to check if it is leetcode solution or not # currently livecodebench only supports LeetCode but # else condition allows future extensibility to other platforms compiled_sol = tmp_sol.Solution() else: # do nothing in the other case since function is accesible compiled_sol = tmp_sol assert compiled_sol is not None finally: signal.alarm(0) return compiled_sol def convert_line_to_decimals(line: str) -> tuple[bool, list[Decimal]]: try: decimal_line = [Decimal(elem) for elem in line.split()] except: return False, [] return True, decimal_line def get_stripped_lines(val: str): ## you don't want empty lines to add empty list after splitlines! val = val.strip() return [val_line.strip() for val_line in val.split("\n")] def grade_call_based( code: str, all_inputs: list, all_outputs: list, fn_name: str, timeout: int ): # call-based clean up logic # need to wrap in try-catch logic after to catch the correct errors, but for now this is fine. code = import_string + "\n\n" + code compiled_sol = compile_code(code, timeout) if compiled_sol is None: return method = get_function(compiled_sol, fn_name) if method is None: return all_inputs = [ [json.loads(line) for line in inputs.split("\n")] for inputs in all_inputs ] all_outputs = [json.loads(output) for output in all_outputs] total_execution = 0 all_results = [] for idx, (gt_inp, gt_out) in enumerate(zip(all_inputs, all_outputs)): signal.alarm(timeout) faulthandler.enable() try: # can lock here so time is useful start = time.time() prediction = method(*gt_inp) total_execution += time.time() - start signal.alarm(0) # don't penalize model if it produces tuples instead of lists # ground truth sequences are not tuples if isinstance(prediction, tuple): prediction = list(prediction) tmp_result = prediction == gt_out # handle floating point comparisons all_results.append(tmp_result) if not tmp_result: return all_results, { "output": truncatefn(prediction), "inputs": truncatefn(gt_inp), "expected": truncatefn(gt_out), "error_code": -2, "error_message": "Wrong Answer", } except Exception as e: signal.alarm(0) if "timeoutexception" in repr(e).lower(): all_results.append(-3) return all_results, { "error": repr(e), "error_code": -3, "error_message": "Time Limit Exceeded", "inputs": truncatefn(gt_inp), "expected": truncatefn(gt_out), } else: all_results.append(-4) return all_results, { "error": repr(e), "error_code": -4, "error_message": "Runtime Error", "inputs": truncatefn(gt_inp), "expected": truncatefn(gt_out), } finally: signal.alarm(0) faulthandler.disable() return all_results, {"execution time": total_execution} def grade_stdio( code: str, all_inputs: list, all_outputs: list, timeout: int, ): ## runtime doesn't interact well with __name__ == '__main__' code = clean_if_name(code) ## we wrap the given code inside another function code = make_function(code) compiled_sol = compile_code(code, timeout) if compiled_sol is None: return method = get_function(compiled_sol, "wrapped_function") if method is None: return all_results = [] total_execution_time = 0 for idx, (gt_inp, gt_out) in enumerate(zip(all_inputs, all_outputs)): signal.alarm(timeout) faulthandler.enable() signal.alarm(timeout) with Capturing() as captured_output: try: start = time.time() call_method(method, gt_inp) total_execution_time += time.time() - start # reset the alarm signal.alarm(0) except Exception as e: signal.alarm(0) if "timeoutexception" in repr(e).lower(): all_results.append(-3) return all_results, { "error": repr(e), "error_code": -3, "error_message": "Time Limit Exceeded", "inputs": truncatefn(gt_inp), "expected": truncatefn(gt_out), } else: all_results.append(-4) return all_results, { "error": repr(e), "error_code": -4, "error_message": "Runtime Error", "inputs": truncatefn(gt_inp), "expected": truncatefn(gt_out), } finally: signal.alarm(0) faulthandler.disable() prediction = captured_output[0] stripped_prediction_lines = get_stripped_lines(prediction) stripped_gt_out_lines = get_stripped_lines(gt_out) ## WA happens in multiple circumstances ## so cache the return to make it clean! WA_send_args = { "output": truncatefn(prediction), "inputs": truncatefn(gt_inp), "expected": truncatefn(gt_out), "error_code": -2, } if len(stripped_prediction_lines) != len(stripped_gt_out_lines): all_results.append(-2) WA_send_args["error_message"] = "Wrong answer: mismatched output length" return all_results, WA_send_args for output_line_idx, ( stripped_prediction_line, stripped_gt_out_line, ) in enumerate(zip(stripped_prediction_lines, stripped_gt_out_lines)): WA_send_args["error_message"] = ( f"Wrong answer at {output_line_idx=}: {truncatefn(stripped_prediction_line)} != {truncatefn(stripped_gt_out_line)}" ) ## CASE 1: exact match if stripped_prediction_line == stripped_gt_out_line: continue ## CASE 2: element-wise comparision ## if there are floating elements ## use `decimal` library for good floating point comparision ## otherwise gotcha: np.isclose(50000000000000000, 50000000000000001) = True ## note that we should always be able to convert to decimals success, decimal_prediction_line = convert_line_to_decimals( stripped_prediction_line ) if not success: all_results.append(-2) return all_results, WA_send_args success, decimal_gtout_line = convert_line_to_decimals(stripped_gt_out_line) if not success: all_results.append(-2) return all_results, WA_send_args if decimal_prediction_line == decimal_gtout_line: continue all_results.append(-2) return all_results, WA_send_args all_results.append(True) return all_results, {"execution time": total_execution_time} def run_test(sample, test=None, debug=False, timeout=6): """ if test(generated_code) is not None it'll try to run the code. otherwise it'll just return an input and output pair. """ signal.signal(signal.SIGALRM, timeout_handler) # Disable functionalities that can make destructive changes to the test. # max memory is set to 4GB reliability_guard() if debug: print(f"start = {datetime.now().time()}") try: in_outs = json.loads(sample["input_output"]) except ValueError as e: raise e in_outs = None if in_outs: if in_outs.get("fn_name") is None: which_type = CODE_TYPE.standard_input # Standard input method_name = None else: which_type = CODE_TYPE.call_based # Call-based method_name = in_outs["fn_name"] if debug: print(f"loaded input_output = {datetime.now().time()}") if test is None: assert False, "should not happen: test code is none" return in_outs, {"error": "no test code provided"} elif test is not None: results = [] sol = import_string if debug: print(f"loading test code = {datetime.now().time()}") if which_type == CODE_TYPE.call_based: signal.alarm(timeout) try: results, metadata = grade_call_based( code=test, all_inputs=in_outs["inputs"], all_outputs=in_outs["outputs"], fn_name=method_name, timeout=timeout, ) return results, metadata except Exception as e: return [-4], { "error_code": -4, "error_message": f"Error during testing: {e}", } finally: signal.alarm(0) elif which_type == CODE_TYPE.standard_input: # sol # if code has if __name__ == "__main__": then remove it signal.alarm(timeout) try: results, metadata = grade_stdio( code=test, all_inputs=in_outs["inputs"], all_outputs=in_outs["outputs"], timeout=timeout, ) return results, metadata except Exception as e: return [-4], { "error_code": -4, "error_message": f"Error during testing: {e}", } finally: signal.alarm(0) def reliability_guard(maximum_memory_bytes=None): """ This disables various destructive functions and prevents the generated code from interfering with the test (e.g. fork bomb, killing other processes, removing filesystem files, etc.) WARNING This function is NOT a security sandbox. Untrusted code, including, model- generated code, should not be blindly executed outside of one. See the Codex paper for more information about OpenAI's code sandbox, and proceed with caution. """ if maximum_memory_bytes is not None: import resource resource.setrlimit( resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes) ) resource.setrlimit( resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes) ) if not platform.uname().system == "Darwin": resource.setrlimit( resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes) ) faulthandler.disable() import builtins # builtins.exit = None builtins.quit = None import os os.environ["OMP_NUM_THREADS"] = "1" os.kill = None os.system = None os.putenv = None os.remove = None os.removedirs = None os.rmdir = None os.fchdir = None os.setuid = None os.fork = None os.forkpty = None os.killpg = None os.rename = None os.renames = None os.truncate = None os.replace = None os.unlink = None os.fchmod = None os.fchown = None os.chmod = None os.chown = None os.chroot = None os.fchdir = None os.lchflags = None os.lchmod = None os.lchown = None os.getcwd = None os.chdir = None import shutil shutil.rmtree = None shutil.move = None shutil.chown = None import subprocess subprocess.Popen = None __builtins__["help"] = None import sys sys.modules["ipdb"] = None sys.modules["joblib"] = None sys.modules["resource"] = None sys.modules["psutil"] = None sys.modules["tkinter"] = None ================================================ FILE: dashboard/dashboard.nix ================================================ { lib , config , dream2nix , ... }: let # Read and parse the lock file rawLockFile = builtins.fromJSON (builtins.readFile "${config.deps.dashboardSrc}/package-lock.json"); # For packages with bundleDependencies, filter out deps that are bundled # (bundled deps are inside the tarball, not separate lockfile entries) fixedPackages = lib.mapAttrs (path: entry: if entry ? bundleDependencies && entry.bundleDependencies != [ ] then entry // { dependencies = lib.filterAttrs (name: _: !(lib.elem name entry.bundleDependencies)) (entry.dependencies or { }); } else entry ) (rawLockFile.packages or { }); fixedLockFile = rawLockFile // { packages = fixedPackages; }; in { imports = [ dream2nix.modules.dream2nix.nodejs-package-lock-v3 dream2nix.modules.dream2nix.nodejs-granular-v3 ]; name = "exo-dashboard"; version = "1.0.0"; mkDerivation = { src = config.deps.dashboardSrc; buildPhase = '' runHook preBuild npm run build runHook postBuild ''; installPhase = '' runHook preInstall cp -r build $out/build runHook postInstall ''; }; deps = { nixpkgs, ... }: { inherit (nixpkgs) stdenv; dashboardSrc = null; # Injected by parts.nix }; nodejs-package-lock-v3 = { # Don't use packageLockFile - provide the fixed lock content directly packageLock = fixedLockFile; }; } ================================================ FILE: dashboard/package.json ================================================ { "name": "exo-dashboard", "private": true, "version": "1.0.0", "type": "module", "scripts": { "dev": "vite dev", "build": "vite build", "preview": "vite preview", "prepare": "svelte-kit sync || echo ''", "check": "svelte-kit sync && svelte-check --tsconfig ./tsconfig.json" }, "devDependencies": { "prettier": "^3.4.2", "prettier-plugin-svelte": "^3.3.3", "@sveltejs/adapter-static": "^3.0.10", "@sveltejs/kit": "^2.48.4", "@sveltejs/vite-plugin-svelte": "^5.0.0", "@tailwindcss/vite": "^4.0.0", "@types/d3": "^7.4.3", "@types/node": "^22", "d3": "^7.9.0", "svelte": "^5.0.0", "svelte-check": "^4.0.0", "tailwindcss": "^4.0.0", "tw-animate-css": "^1.3.5", "typescript": "^5.0.0", "vite": "^6.0.0" }, "dependencies": { "highlight.js": "^11.11.1", "katex": "^0.16.27", "marked": "^17.0.1", "mode-watcher": "^1.1.0" } } ================================================ FILE: dashboard/parts.nix ================================================ { inputs, ... }: { perSystem = { pkgs, lib, ... }: let # Filter source to ONLY include package.json and package-lock.json # This ensures prettier-svelte only rebuilds when lockfiles change dashboardLockfileSrc = lib.cleanSourceWith { src = inputs.self; filter = path: type: let baseName = builtins.baseNameOf path; isDashboardDir = baseName == "dashboard" && type == "directory"; isPackageFile = (lib.hasInfix "/dashboard/" path || lib.hasSuffix "/dashboard" (builtins.dirOf path)) && (baseName == "package.json" || baseName == "package-lock.json"); in isDashboardDir || isPackageFile; }; # Stub source with lockfiles and minimal files for build to succeed # This allows prettier-svelte to avoid rebuilding when dashboard source changes dashboardStubSrc = pkgs.runCommand "dashboard-stub-src" { } '' mkdir -p $out cp ${dashboardLockfileSrc}/dashboard/package.json $out/ cp ${dashboardLockfileSrc}/dashboard/package-lock.json $out/ # Minimal files so vite build succeeds (produces empty output) echo '' > $out/index.html mkdir -p $out/src touch $out/src/app.html ''; # Deps-only build using stub source (for prettier-svelte) # Only rebuilds when package.json or package-lock.json change dashboardDeps = inputs.dream2nix.lib.evalModules { packageSets.nixpkgs = pkgs; modules = [ ./dashboard.nix { paths.projectRoot = inputs.self; paths.projectRootFile = "flake.nix"; paths.package = inputs.self + "/dashboard"; } { deps.dashboardSrc = lib.mkForce dashboardStubSrc; } # Override build phases to skip the actual build - just need node_modules { mkDerivation = { buildPhase = lib.mkForce "true"; installPhase = lib.mkForce '' runHook preInstall runHook postInstall ''; }; } ]; }; # Filter source to only include dashboard directory dashboardSrc = lib.cleanSourceWith { src = inputs.self; filter = path: type: let baseName = builtins.baseNameOf path; inDashboardDir = (lib.hasInfix "/dashboard/" path) || (lib.hasSuffix "/dashboard" (builtins.dirOf path)) || (baseName == "dashboard" && type == "directory"); in inDashboardDir; }; # Build the dashboard with dream2nix (includes node_modules in output) dashboardFull = inputs.dream2nix.lib.evalModules { packageSets.nixpkgs = pkgs; modules = [ ./dashboard.nix { paths.projectRoot = inputs.self; paths.projectRootFile = "flake.nix"; paths.package = inputs.self + "/dashboard"; } # Inject the filtered source { deps.dashboardSrc = lib.mkForce "${dashboardSrc}/dashboard"; } ]; }; in { # Extract just the static site from the full build packages.dashboard = pkgs.runCommand "exo-dashboard" { } '' cp -r ${dashboardFull}/build $out ''; # Prettier with svelte plugin for treefmt # Uses dashboardDeps instead of dashboardFull to avoid rebuilding on source changes packages.prettier-svelte = pkgs.writeShellScriptBin "prettier-svelte" '' export NODE_PATH="${dashboardDeps}/lib/node_modules/exo-dashboard/node_modules" exec ${pkgs.nodejs}/bin/node \ ${dashboardDeps}/lib/node_modules/exo-dashboard/node_modules/prettier/bin/prettier.cjs \ --plugin "${dashboardDeps}/lib/node_modules/exo-dashboard/node_modules/prettier-plugin-svelte/plugin.js" \ "$@" ''; }; } ================================================ FILE: dashboard/src/app.css ================================================ @import 'tailwindcss'; @import 'tw-animate-css'; @custom-variant dark (&:is(.dark *)); :root { /* EXO Brand Colors - Command Center Theme (neutral dark greys) */ --exo-black: oklch(0.12 0 0); --exo-dark-gray: oklch(0.16 0 0); --exo-medium-gray: oklch(0.22 0 0); --exo-light-gray: oklch(0.6 0 0); --exo-yellow: oklch(0.85 0.18 85); --exo-yellow-darker: oklch(0.7 0.16 85); --exo-yellow-glow: oklch(0.9 0.2 85); /* Gotham-inspired accent colors */ --exo-grid: oklch(0.25 0 0); --exo-scanline: oklch(0.15 0 0); --exo-glow-yellow: 0 0 20px oklch(0.85 0.18 85 / 0.3); --exo-glow-yellow-strong: 0 0 40px oklch(0.85 0.18 85 / 0.5); /* Theme Variables */ --radius: 0.375rem; --background: var(--exo-black); --foreground: oklch(0.9 0 0); --card: var(--exo-dark-gray); --card-foreground: oklch(0.9 0 0); --popover: var(--exo-dark-gray); --popover-foreground: oklch(0.9 0 0); --primary: var(--exo-yellow); --primary-foreground: var(--exo-black); --secondary: var(--exo-medium-gray); --secondary-foreground: oklch(0.9 0 0); --muted: var(--exo-medium-gray); --muted-foreground: var(--exo-light-gray); --accent: var(--exo-medium-gray); --accent-foreground: oklch(0.9 0 0); --destructive: oklch(0.6 0.25 25); --border: oklch(0.22 0 0); --input: oklch(0.22 0 0); --ring: var(--exo-yellow); } @theme inline { --radius-sm: calc(var(--radius) - 2px); --radius-md: var(--radius); --radius-lg: calc(var(--radius) + 2px); --radius-xl: calc(var(--radius) + 4px); --color-background: var(--background); --color-foreground: var(--foreground); --color-card: var(--card); --color-card-foreground: var(--card-foreground); --color-popover: var(--popover); --color-popover-foreground: var(--popover-foreground); --color-primary: var(--primary); --color-primary-foreground: var(--primary-foreground); --color-secondary: var(--secondary); --color-secondary-foreground: var(--secondary-foreground); --color-muted: var(--muted); --color-muted-foreground: var(--muted-foreground); --color-accent: var(--accent); --color-accent-foreground: var(--accent-foreground); --color-destructive: var(--destructive); --color-border: var(--border); --color-input: var(--input); --color-ring: var(--ring); /* Custom EXO colors */ --color-exo-yellow: var(--exo-yellow); --color-exo-yellow-darker: var(--exo-yellow-darker); --color-exo-black: var(--exo-black); --color-exo-dark-gray: var(--exo-dark-gray); --color-exo-medium-gray: var(--exo-medium-gray); --color-exo-light-gray: var(--exo-light-gray); } @layer base { * { @apply border-border outline-ring/50; } html, body { @apply bg-background text-foreground; font-family: 'SF Mono', 'Fira Code', 'Monaco', 'Consolas', 'Liberation Mono', monospace; letter-spacing: 0.02em; } } @layer utilities { .scrollbar-hide { &::-webkit-scrollbar { display: none; } -ms-overflow-style: none; scrollbar-width: none; } /* CRT Scanline effect */ .scanlines { position: relative; &::before { content: ''; position: absolute; inset: 0; background: repeating-linear-gradient( 0deg, transparent, transparent 2px, oklch(0 0 0 / 0.03) 2px, oklch(0 0 0 / 0.03) 4px ); pointer-events: none; z-index: 100; } } /* Command panel styling */ .command-panel { background: linear-gradient( 180deg, oklch(0.16 0 0 / 0.95) 0%, oklch(0.12 0 0 / 0.98) 100% ); border: 1px solid oklch(0.25 0 0); box-shadow: inset 0 1px 0 oklch(1 0 0 / 0.03), 0 4px 20px oklch(0 0 0 / 0.5); } /* Glow text */ .glow-text { text-shadow: 0 0 10px oklch(0.85 0.18 85 / 0.5), 0 0 20px oklch(0.85 0.18 85 / 0.3), 0 0 40px oklch(0.85 0.18 85 / 0.1); } /* Status indicator pulse */ .status-pulse { animation: statusPulse 2s ease-in-out infinite; } /* Grid background */ .grid-bg { background-image: linear-gradient(oklch(0.2 0 0 / 0.3) 1px, transparent 1px), linear-gradient(90deg, oklch(0.2 0 0 / 0.3) 1px, transparent 1px); background-size: 40px 40px; } } /* Animations */ @keyframes flowAnimation { from { stroke-dashoffset: 0; } to { stroke-dashoffset: -16; } } @keyframes statusPulse { 0%, 100% { opacity: 1; } 50% { opacity: 0.5; } } @keyframes radarSweep { from { transform: rotate(0deg); } to { transform: rotate(360deg); } } @keyframes glowPulse { 0%, 100% { box-shadow: 0 0 5px oklch(0.85 0.18 85 / 0.3), 0 0 10px oklch(0.85 0.18 85 / 0.1); } 50% { box-shadow: 0 0 15px oklch(0.85 0.18 85 / 0.5), 0 0 30px oklch(0.85 0.18 85 / 0.2); } } @keyframes dataPulse { 0%, 100% { opacity: 0.6; } 50% { opacity: 1; } } .graph-link { stroke: oklch(0.85 0.18 85 / 0.4); stroke-width: 1.5px; stroke-dasharray: 8, 8; animation: flowAnimation 1s linear infinite; filter: drop-shadow(0 0 3px oklch(0.85 0.18 85 / 0.5)); } /* Onboarding step 2: connection line between devices */ .onboarding-connection-line { stroke: oklch(0.85 0.18 85 / 0.5); stroke-width: 1.5px; stroke-dasharray: 6, 6; animation: flowAnimation 1s linear infinite; filter: drop-shadow(0 0 4px oklch(0.85 0.18 85 / 0.4)); } /* Onboarding step 4: red connection line for disconnect */ .onboarding-connection-line-red { stroke: rgba(220, 38, 38, 0.7); stroke-width: 1.5px; stroke-dasharray: 6, 6; filter: drop-shadow(0 0 2px rgba(220, 38, 38, 0.3)); } .graph-link-active { stroke: oklch(0.85 0.18 85 / 0.8); stroke-width: 2px; filter: drop-shadow(0 0 6px oklch(0.85 0.18 85 / 0.8)); } /* CRT Screen effect for topology */ .crt-screen { position: relative; border-radius: 50%; background: radial-gradient( ellipse at center, oklch(0.16 0 0) 0%, oklch(0.12 0 0) 50%, oklch(0.09 0 0) 100% ); box-shadow: inset 0 0 100px oklch(0 0 0 / 0.5), 0 0 50px oklch(0.85 0.18 85 / 0.1); } /* Data readout styling */ .data-readout { font-family: 'SF Mono', 'Fira Code', monospace; font-size: 11px; letter-spacing: 0.05em; text-transform: uppercase; } /* Terminal cursor blink */ .cursor-blink { animation: cursorBlink 1s step-end infinite; } @keyframes cursorBlink { 0%, 100% { opacity: 1; } 50% { opacity: 0; } } /* Custom scrollbar for command center */ ::-webkit-scrollbar { width: 6px; height: 6px; } ::-webkit-scrollbar-track { background: oklch(0.1 0 0); } ::-webkit-scrollbar-thumb { background: oklch(0.3 0 0); border-radius: 3px; } ::-webkit-scrollbar-thumb:hover { background: oklch(0.85 0.18 85 / 0.5); } /* Remove focus outline/border for inputs */ input:focus, textarea:focus { outline: none; box-shadow: none; } /* Shooting Stars Animation */ .shooting-stars { position: fixed; inset: 0; overflow: hidden; pointer-events: none; z-index: 0; } .shooting-star { position: absolute; width: 3px; height: 3px; background: oklch(0.85 0.18 85 / 1); border-radius: 50%; box-shadow: 0 0 6px oklch(0.85 0.18 85 / 0.8); animation: shootingStar var(--duration, 3s) linear infinite; animation-delay: var(--delay, 0s); opacity: 0; } .shooting-star::before { content: ''; position: absolute; width: 80px; height: 2px; background: linear-gradient(90deg, oklch(0.85 0.18 85 / 0), oklch(0.85 0.18 85 / 0.6)); transform: rotate(45deg); transform-origin: right center; top: 0; right: 2px; } @keyframes shootingStar { 0% { opacity: 0; transform: translate(0, 0); } 0.5% { opacity: 1; } 2.5% { opacity: 0.8; transform: translate(300px, 300px); } 3.5% { opacity: 0; transform: translate(400px, 400px); } 100% { opacity: 0; transform: translate(400px, 400px); } } /* Onboarding smooth fade-in */ @keyframes onb-fade-in { from { opacity: 0; transform: translateY(8px); } to { opacity: 1; transform: none; } } /* Opacity-only fade-in — no transform, safe for containers with fixed-position children */ @keyframes onb-fade-opacity { from { opacity: 0; } to { opacity: 1; } } /* Respect reduced motion preference */ @media (prefers-reduced-motion: reduce) { .shooting-star, .shooting-star::before { animation: none !important; opacity: 0 !important; } .graph-link { animation: none; } .status-pulse { animation: none; } .cursor-blink { animation: none; } .onboarding-connection-line { animation: none; } .onboarding-connection-line-red { animation: none; } [style*="onb-fade-in"], [style*="onb-fade-opacity"] { animation: none !important; opacity: 1 !important; } *, *::before, *::after { transition-duration: 0.01ms !important; animation-duration: 0.01ms !important; animation-iteration-count: 1 !important; } } ================================================ FILE: dashboard/src/app.d.ts ================================================ // See https://svelte.dev/docs/kit/types#app.d.ts // for information about these interfaces declare global { namespace App { // interface Error {} // interface Locals {} // interface PageData {} // interface PageState {} // interface Platform {} } } export {}; ================================================ FILE: dashboard/src/app.html ================================================ EXO %sveltekit.head%
%sveltekit.body%
================================================ FILE: dashboard/src/lib/components/ChatAttachments.svelte ================================================ {#if files.length > 0}
{#each files as file (file.id)}
{#if file.preview && getFileCategory(file.type, file.name) === "image"} {file.name} {:else} {getFileIcon(file)} {/if}
{truncateName(file.name)} {formatFileSize(file.size)}
{#if !readonly && onRemove} {/if}
{/each}
{/if} ================================================ FILE: dashboard/src/lib/components/ChatForm.svelte ================================================
{ e.preventDefault(); handleSubmit(); }} class="w-full {className}" ondragover={handleDragOver} ondragleave={handleDragLeave} ondrop={handleDrop} >
{#if isDragOver}
DROP FILES HERE
{/if} {#if isEditMode && currentEditingImage}
Source for editing
EDITING IMAGE
{/if} {#if showModelSelector}
MODEL:
{#if modelSupportsThinking()} {/if} {#if currentTtft !== null || currentTps !== null}
{#if currentTtft !== null} TTFT {currentTtft.toFixed(1)}ms {/if} {#if currentTps !== null} TPS {currentTps.toFixed(1)} tok/s ({(1000 / currentTps).toFixed(1)} ms/tok) {/if}
{/if}
{/if} {#if showModelSelector && (isImageModel() || isEditMode)} {/if} {#if uploadedFiles.length > 0}
{/if}
▶ {#if loading} {:else} {/if}
{#if showHelperText}

ENTER TO SEND | SHIFT+ENTER NEW LINE | DRAG & DROP OR PASTE FILES

{/if}
================================================ FILE: dashboard/src/lib/components/ChatMessages.svelte ================================================
{#each messageList as message, i (message.id)}
{#if message.role === "assistant"}
EXO {formatTimestamp(message.timestamp)} {#if message.ttftMs || message.tps} {#if message.ttftMs}TTFT {message.ttftMs.toFixed( 0, )}ms{/if}{#if message.ttftMs && message.tps}•{/if}{#if message.tps}{message.tps.toFixed(1)} tok/s{/if} {/if}
{:else}
{formatTimestamp(message.timestamp)} QUERY
{/if} {#if deleteConfirmId === message.id}

{#if i === messageList.length - 1} Delete this message? {:else} Delete this message and all messages after it? {/if}

{:else if editingMessageId === message.id}
{:else}
{#if message.role === "user"}
{#if message.attachments && message.attachments.length > 0}
{#each message.attachments as attachment}
{#if attachment.type === "image" && attachment.preview} {attachment.name} { if (attachment.preview) expandedImageSrc = attachment.preview; }} /> {:else} {getAttachmentIcon(attachment)} {/if} {truncateName(attachment.name)}
{/each}
{/if} {#if message.content}
{message.content}
{/if}
{:else}
{#if loading && isLastAssistantMessage(message.id) && prefill && !message.content} {/if} {#if message.thinking && message.thinking.trim().length > 0}
{#if isThinkingExpanded(message.id)}
{message.thinking.trim()}
{/if}
{/if} {#if message.attachments?.some((a) => a.type === "generated-image")}
{#each message.attachments.filter((a) => a.type === "generated-image") as attachment}
{ if (attachment.preview) expandedImageSrc = attachment.preview; }} />
{/each}
{/if}
{#if message.content === "Generating image..." || message.content === "Editing image..." || message.content?.startsWith("Generating...") || message.content?.startsWith("Editing...")}
{message.content}
{:else if message.content || (loading && !message.attachments?.some((a) => a.type === "generated-image"))} {#if isHeatmapVisible(message.id) && message.tokens && message.tokens.length > 0} regenerateFromToken(message.id, tokenIndex)} /> {:else} {#if loading && !message.content} {/if} {/if} {/if}
{/if}
{#if message.role === "user"} {/if} {#if message.role === "assistant" && message.tokens && message.tokens.length > 0} {/if} {#if message.role === "assistant" && isLastAssistantMessage(message.id) && !loading} {/if}
{/if}
{/each} {#if messageList.length === 0}

AWAITING INPUT

Type a message below · Shift+Enter for newline

{/if}
{#if showScrollButton} {/if}
(expandedImageSrc = null)} /> ================================================ FILE: dashboard/src/lib/components/ChatModelSelector.svelte ================================================

Recommended for your

{clusterLabel}

{#each recommendations as rec} {#if rec.model} {:else}
{rec.label}

No model fits

{/if} {/each}

Or just start typing — we'll pick the best model automatically

{#if hoveredTooltip && tooltipAnchor} {@const rec = recommendations.find((r) => r.category === hoveredTooltip)} {#if rec}
{rec.tooltip}
{/if} {/if} ================================================ FILE: dashboard/src/lib/components/ChatSidebar.svelte ================================================ {#snippet sidebarContent()}
{#if filteredConversations.length > 0}
{searchQuery ? "SEARCH RESULTS" : "CONVERSATIONS"}
{#each filteredConversations as conversation (conversation.id)} {@const info = resolveConversationInfo(conversation)}
{#if editingId === conversation.id}
{:else if deleteConfirmId === conversation.id}

Delete "{conversation.name}"?

{:else} {@const stats = getLastAssistantStats(conversation)}
handleSelectConversation(conversation.id)} onkeydown={(e) => e.key === "Enter" && handleSelectConversation(conversation.id)} class="group w-full flex items-center justify-between p-2.5 rounded-lg mb-1 transition-all text-left cursor-pointer {activeId === conversation.id ? 'bg-exo-yellow/5 border border-exo-yellow/30' : 'hover:bg-white/[0.03] hover:border-white/10 border border-transparent'}" >
{conversation.name}
{formatDate(conversation.updatedAt)}
{info.modelLabel}
{#if stats}
{#if stats.ttftMs}TTFT {stats.ttftMs.toFixed(0)}ms{/if}{#if stats.ttftMs && stats.tps}·{/if}{#if stats.tps}{stats.tps.toFixed(1)} tok/s{/if}
{/if}
{/if}
{/each}
{:else}

{searchQuery ? "NO RESULTS" : "NO CONVERSATIONS"}

{searchQuery ? "Try a different search" : "Start a new chat to begin"}

{/if}
{#if showDeleteAllConfirm}

Delete all {conversationList.length} conversations?

{:else if conversationList.length > 0} {/if}
{conversationList.length} CONVERSATION{conversationList.length !== 1 ? "S" : ""}
{/snippet} {#if isMobileDrawer} {#if isOpen} {/if} {:else} {/if} ================================================ FILE: dashboard/src/lib/components/ConnectionBanner.svelte ================================================ {#if !connected} {/if} ================================================ FILE: dashboard/src/lib/components/DeviceIcon.svelte ================================================ {#if modelLower === "mac studio" || modelLower === "mac mini"} {#if ramPercent > 0} {/if} {:else} {#if ramPercent > 0} {/if} {/if} ================================================ FILE: dashboard/src/lib/components/FamilyLogos.svelte ================================================ {#if family === "favorites"} {:else if family === "recents"} {:else if family === "llama" || family === "meta"} {:else if family === "qwen"} {:else if family === "deepseek"} {:else if family === "openai" || family === "gpt-oss"} {:else if family === "glm"} {:else if family === "minimax"} {:else if family === "kimi"} {:else if family === "flux"} {:else if family === "qwen-image"} {:else if family === "huggingface"} {:else if family === "step"} {:else} {/if} ================================================ FILE: dashboard/src/lib/components/FamilySidebar.svelte ================================================
{#if hasFavorites} {/if} {#if hasRecents} {/if}
{#each families as family} {/each}
================================================ FILE: dashboard/src/lib/components/HeaderNav.svelte ================================================
================================================ FILE: dashboard/src/lib/components/HuggingFaceResultItem.svelte ================================================
{modelName} {#if downloadedOnNodes.length > 0} {/if} {#if isAdded} Added {/if}
{model.author} {formatNumber(model.downloads)} {formatNumber(model.likes)}
{#if isAdded} {:else} {/if}
================================================ FILE: dashboard/src/lib/components/ImageLightbox.svelte ================================================ {#if src} {/if} ================================================ FILE: dashboard/src/lib/components/ImageParamsPanel.svelte ================================================
SIZE:
{#if isSizeDropdownOpen}
{#each sizeOptions as size} {/each}
{/if}
QUALITY:
{#if isQualityDropdownOpen}
{#each qualityOptions as quality} {/each}
{/if}
FORMAT:
{#each outputFormatOptions as format} {/each}
{#if !isEditMode}
IMAGES:
{/if}
STREAM:
{#if params.stream}
PARTIALS:
{/if} {#if isEditMode}
FIDELITY:
{#each inputFidelityOptions as fidelity} {/each}
{/if}
{#if showAdvanced}
SEED:
STEPS:
{params.numInferenceSteps ?? "--"} {#if params.numInferenceSteps !== null} {/if}
GUIDANCE:
{params.guidance !== null ? params.guidance.toFixed(1) : "--"} {#if params.guidance !== null} {/if}
SYNC STEPS:
{params.numSyncSteps ?? "--"} {#if params.numSyncSteps !== null} {/if}
NEGATIVE PROMPT:
{/if}
================================================ FILE: dashboard/src/lib/components/MarkdownContent.svelte ================================================
{@html processedHtml}
================================================ FILE: dashboard/src/lib/components/ModelCard.svelte ================================================
{model.name || model.id}
{#if huggingFaceModelId} {/if} {#if tags.length > 0}
{#each tags as tag} {tag} {/each}
{/if}
{#if model.name && model.name !== model.id}
{model.id}
{/if}
{estimatedMemory}GB
{sharding} {runtime === "MlxRing" ? "MLX Ring" : runtime === "MlxJaccl" ? "MLX RDMA" : runtime}
{#if perNode.length > 0}
Download progress
{#each perNode as node}
{node.nodeName}
{#if node.status === "downloading" && node.progress} {Math.round(node.percentage)}% {formatSpeed( node.progress.speed, )} {:else} {node.percentage > 0 ? `${Math.round(node.percentage)}%` : "0%"} {/if}
{/each}
{/if} {#if placementPreview().nodes.length > 0} {@const preview = placementPreview()}
{#if preview.nodes.length > 1} {@const usedNodes = preview.nodes.filter((n) => n.isUsed)} {@const nodePositions = Object.fromEntries( preview.nodes.map((n) => [n.id, { x: n.x, y: n.y }]), )} {@const allConnections = isDebugMode && usedNodes.length > 1 ? (() => { const conns: Array = []; for (let i = 0; i < usedNodes.length; i++) { for (let j = i + 1; j < usedNodes.length; j++) { const n1 = usedNodes[i]; const n2 = usedNodes[j]; const midX = (n1.x + n2.x) / 2; const midY = (n1.y + n2.y) / 2; for (const c of getConnectionInfo(n1.id, n2.id)) { const fromPos = nodePositions[c.from]; const toPos = nodePositions[c.to]; const arrow = fromPos && toPos ? getArrow(fromPos, toPos) : "→"; conns.push({ ...c, midX, midY, arrow, }); } } } return conns; })() : []} {#each preview.nodes as node, i} {#each preview.nodes.slice(i + 1) as node2} {/each} {/each} {#if isDebugMode && allConnections.length > 0} {@const centerX = preview.topoWidth / 2} {@const centerY = preview.topoHeight / 2} {@const quadrants = { topLeft: allConnections.filter( (c) => c.midX < centerX && c.midY < centerY, ), topRight: allConnections.filter( (c) => c.midX >= centerX && c.midY < centerY, ), bottomLeft: allConnections.filter( (c) => c.midX < centerX && c.midY >= centerY, ), bottomRight: allConnections.filter( (c) => c.midX >= centerX && c.midY >= centerY, ), }} {@const padding = 4} {@const lineHeight = 8} {#each quadrants.topLeft as conn, idx} {conn.arrow} {isRdma ? conn.iface || "?" : `${conn.ip}${conn.iface ? ` (${conn.iface})` : ""}`} {/each} {#each quadrants.topRight as conn, idx} {conn.arrow} {isRdma ? conn.iface || "?" : `${conn.ip}${conn.iface ? ` (${conn.iface})` : ""}`} {/each} {#each quadrants.bottomLeft as conn, idx} {conn.arrow} {isRdma ? conn.iface || "?" : `${conn.ip}${conn.iface ? ` (${conn.iface})` : ""}`} {/each} {#each quadrants.bottomRight as conn, idx} {conn.arrow} {isRdma ? conn.iface || "?" : `${conn.ip}${conn.iface ? ` (${conn.iface})` : ""}`} {/each} {/if} {/if} {#each preview.nodes as node} {#if node.deviceType === "macbook"} {#if node.modelUsageGB > 0 && node.isUsed} {/if} {:else if node.deviceType === "studio"} {#if node.modelUsageGB > 0 && node.isUsed} {/if} {:else if node.deviceType === "mini"} {#if node.modelUsageGB > 0 && node.isUsed} {/if} {:else} {/if} 90 ? "#f87171" : "#FFD700" : "#4B5563"} > {node.newPercent.toFixed(0)}% {/each}
{/if}
================================================ FILE: dashboard/src/lib/components/ModelFilterPopover.svelte ================================================
e.stopPropagation()} role="dialog" aria-label="Filter options" >

Capabilities

{#each availableCapabilities as cap} {@const isSelected = filters.capabilities.includes(cap.id)} {/each}

Availability

Model Size

{#each sizeRanges as range} {@const isSelected = filters.sizeRange && filters.sizeRange.min === range.min && filters.sizeRange.max === range.max} {/each}
================================================ FILE: dashboard/src/lib/components/ModelPickerGroup.svelte ================================================
v.id).join(" ")} class="border-b border-white/5 last:border-b-0 {!anyVariantFits && !anyVariantHasInstance ? 'opacity-50' : ''} {isHighlighted ? 'model-just-added' : ''}" >
{ if (group.hasMultipleVariants) { onToggleExpand(); } else { const modelId = group.variants[0]?.id; if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) { onSelectModel(modelId); } } }} role="button" tabindex="0" onkeydown={(e) => { if (e.key === "Enter" || e.key === " ") { e.preventDefault(); if (group.hasMultipleVariants) { onToggleExpand(); } else { const modelId = group.variants[0]?.id; if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) { onSelectModel(modelId); } } } }} > {#if group.hasMultipleVariants} {:else}
{/if}
{group.name} {#each group.capabilities.filter((c) => c !== "text") as cap} {#if cap === "thinking"} {:else if cap === "code"} {:else if cap === "vision"} {:else if cap === "image_gen"} {:else if cap === "image_edit"} {/if} {/each}
{#if !group.hasMultipleVariants && group.smallestVariant?.storage_size_megabytes} {@const singleVariantFitStatus = getModelFitStatus( group.smallestVariant.id, )} {formatSize(group.smallestVariant.storage_size_megabytes)} {/if} {#if group.hasMultipleVariants} {@const sizes = group.variants .map((v) => v.storage_size_megabytes || 0) .filter((s) => s > 0) .sort((a, b) => a - b)} {group.variants.length} variants{#if sizes.length >= 2}{" "}({formatSize( sizes[0], )}-{formatSize(sizes[sizes.length - 1])}){/if} {/if} {#if launchedAt} {timeAgo(launchedAt)} {/if} {#if groupDownloadStatus && groupDownloadStatus.nodeIds.length > 0} {/if} {#if groupInstanceStatus} {#if groupInstanceStatus.status === "READY" || groupInstanceStatus.status === "LOADED" || groupInstanceStatus.status === "RUNNING"} {:else if groupInstanceStatus.status === "DOWNLOADING"} {:else if groupInstanceStatus.status === "LOADING" || groupInstanceStatus.status === "WARMING UP"} {/if} {/if} {#if isMainSelected} {/if}
{#if isExpanded && group.hasMultipleVariants}
{#each group.variants as variant} {@const fitStatus = getModelFitStatus(variant.id)} {@const modelCanFit = canModelFit(variant.id)} {@const variantHasInstance = instanceStatuses[variant.id] != null} {@const isSelected = selectedModelId === variant.id}
{ if (modelCanFit || variantHasInstance) { onSelectModel(variant.id); } }} onkeydown={(e) => { if (e.key === "Enter" || e.key === " ") { e.preventDefault(); if (modelCanFit) { onSelectModel(variant.id); } } }} > {variant.quantization || "default"} {formatSize(variant.storage_size_megabytes)} {#if downloadStatusMap?.get(variant.id)} {@const variantDl = downloadStatusMap.get(variant.id)} {#if variantDl} {/if} {/if} {#if instanceStatuses[variant.id]} {@const instStatus = instanceStatuses[variant.id]} {#if instStatus.status === "READY" || instStatus.status === "LOADED" || instStatus.status === "RUNNING"} {:else if instStatus.status === "DOWNLOADING"} {:else if instStatus.status === "LOADING" || instStatus.status === "WARMING UP"} {/if} {/if} {#if isSelected} {/if}
{/each}
{/if}
================================================ FILE: dashboard/src/lib/components/ModelPickerModal.svelte ================================================ {#if isOpen} {#if infoGroup}
(infoGroup = null)} role="presentation" >
{/if} {/if} ================================================ FILE: dashboard/src/lib/components/PrefillProgressBar.svelte ================================================
Processing prompt {formatTokenCount(progress.processed)} / {formatTokenCount( progress.total, )} tokens
{etaText ?? ""} {percentage}%
================================================ FILE: dashboard/src/lib/components/ToastContainer.svelte ================================================ {#if items.length > 0}
{#each items as toast (toast.id)} {@const style = typeStyles[toast.type]} {/each}
{/if} ================================================ FILE: dashboard/src/lib/components/TokenHeatmap.svelte ================================================
{#each tokens as tokenData, i (i)} handleMouseEnter(e, tokenData, i)} onmouseleave={handleMouseLeave}>{tokenData.token} {/each}
{#if hoveredToken}
Token: "{hoveredToken.token.token}" {formatProbability(hoveredToken.token.probability)}
logprob: {formatLogprob(hoveredToken.token.logprob)}
{#if hoveredToken.token.topLogprobs.length > 0}
Alternatives:
{#each hoveredToken.token.topLogprobs.slice(0, 5) as alt, idx (idx)} {@const altProb = Math.exp(alt.logprob)}
"{alt.token}" {formatProbability(altProb)}
{/each}
{/if} {#if onRegenerateFrom} {/if}
{/if} ================================================ FILE: dashboard/src/lib/components/TopologyGraph.svelte ================================================ ================================================ FILE: dashboard/src/lib/components/index.ts ================================================ export { default as TopologyGraph } from "./TopologyGraph.svelte"; export { default as ChatForm } from "./ChatForm.svelte"; export { default as ChatMessages } from "./ChatMessages.svelte"; export { default as ChatAttachments } from "./ChatAttachments.svelte"; export { default as ChatSidebar } from "./ChatSidebar.svelte"; export { default as ModelCard } from "./ModelCard.svelte"; export { default as MarkdownContent } from "./MarkdownContent.svelte"; export { default as ImageParamsPanel } from "./ImageParamsPanel.svelte"; export { default as FamilyLogos } from "./FamilyLogos.svelte"; export { default as FamilySidebar } from "./FamilySidebar.svelte"; export { default as HuggingFaceResultItem } from "./HuggingFaceResultItem.svelte"; export { default as ModelFilterPopover } from "./ModelFilterPopover.svelte"; export { default as ModelPickerGroup } from "./ModelPickerGroup.svelte"; export { default as ModelPickerModal } from "./ModelPickerModal.svelte"; export { default as ChatModelSelector } from "./ChatModelSelector.svelte"; ================================================ FILE: dashboard/src/lib/stores/app.svelte.ts ================================================ /** * AppStore - Central state management for the EXO dashboard * * Manages: * - Chat state (whether a conversation has started) * - Topology data from the EXO server * - UI state for the topology/chat transition */ import { browser } from "$app/environment"; // UUID generation fallback for browsers without crypto.randomUUID function generateUUID(): string { if ( typeof crypto !== "undefined" && typeof crypto.randomUUID === "function" ) { return crypto.randomUUID(); } // Fallback implementation return "xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx".replace(/[xy]/g, (c) => { const r = (Math.random() * 16) | 0; const v = c === "x" ? r : (r & 0x3) | 0x8; return v.toString(16); }); } export interface NodeInfo { system_info?: { model_id?: string; chip?: string; memory?: number; }; network_interfaces?: Array<{ name?: string; addresses?: string[]; }>; ip_to_interface?: Record; macmon_info?: { memory?: { ram_usage: number; ram_total: number; }; temp?: { gpu_temp_avg: number; }; gpu_usage?: [number, number]; sys_power?: number; }; last_macmon_update: number; friendly_name?: string; os_version?: string; } export interface TopologyEdge { source: string; target: string; sendBackIp?: string; sendBackInterface?: string; sourceRdmaIface?: string; sinkRdmaIface?: string; } export interface TopologyData { nodes: Record; edges: TopologyEdge[]; } export interface Instance { shardAssignments?: { modelId?: string; runnerToShard?: Record; nodeToRunner?: Record; }; } // Granular node state types from the new state structure interface RawNodeIdentity { modelId?: string; chipId?: string; friendlyName?: string; osVersion?: string; osBuildVersion?: string; } interface RawMemoryUsage { ramTotal?: { inBytes: number }; ramAvailable?: { inBytes: number }; swapTotal?: { inBytes: number }; swapAvailable?: { inBytes: number }; } interface RawSystemPerformanceProfile { gpuUsage?: number; temp?: number; sysPower?: number; pcpuUsage?: number; ecpuUsage?: number; } interface RawNetworkInterfaceInfo { name?: string; ipAddress?: string; addresses?: Array<{ address?: string } | string>; ipv4?: string; ipv6?: string; ipAddresses?: string[]; ips?: string[]; } interface RawNodeNetworkInfo { interfaces?: RawNetworkInterfaceInfo[]; } interface RawSocketConnection { sinkMultiaddr?: { address?: string; ip_address?: string; address_type?: string; port?: number; }; } interface RawRDMAConnection { sourceRdmaIface?: string; sinkRdmaIface?: string; } type RawConnectionEdge = RawSocketConnection | RawRDMAConnection; // New nested mapping format: { source: { sink: [edge1, edge2, ...] } } type RawConnectionsMap = Record>; interface RawTopology { nodes: string[]; connections?: RawConnectionsMap; } export interface DownloadProgress { totalBytes: number; downloadedBytes: number; speed: number; etaMs: number; percentage: number; completedFiles: number; totalFiles: number; files: Array<{ name: string; totalBytes: number; downloadedBytes: number; speed: number; etaMs: number; percentage: number; }>; } export interface ModelDownloadStatus { isDownloading: boolean; progress: DownloadProgress | null; nodeDetails: Array<{ nodeId: string; nodeName: string; progress: DownloadProgress; }>; } // Placement preview from the API export interface PlacementPreview { model_id: string; sharding: "Pipeline" | "Tensor"; instance_meta: "MlxRing" | "MlxJaccl"; instance: unknown | null; memory_delta_by_node: Record | null; error: string | null; } export interface PlacementPreviewResponse { previews: PlacementPreview[]; } interface ImageApiResponse { created: number; data: Array<{ b64_json?: string; url?: string }>; } // Trace API response types export interface TraceCategoryStats { totalUs: number; count: number; minUs: number; maxUs: number; avgUs: number; } export interface TraceRankStats { byCategory: Record; } export interface TraceStatsResponse { taskId: string; totalWallTimeUs: number; byCategory: Record; byRank: Record; } export interface TraceListItem { taskId: string; createdAt: string; fileSize: number; } export interface TraceListResponse { traces: TraceListItem[]; } interface RawStateResponse { topology?: RawTopology; instances?: Record< string, { MlxRingInstance?: Instance; MlxJacclInstance?: Instance; } >; runners?: Record; downloads?: Record; // New granular node state fields nodeIdentities?: Record; nodeMemory?: Record; nodeSystem?: Record; nodeNetwork?: Record; // Thunderbolt identifiers per node nodeThunderbolt?: Record< string, { interfaces: Array<{ rdmaInterface: string; domainUuid: string; linkSpeed: string; }>; } >; // RDMA ctl status per node nodeRdmaCtl?: Record; // Thunderbolt bridge status per node nodeThunderboltBridge?: Record< string, { enabled: boolean; exists: boolean; serviceName?: string | null } >; // Thunderbolt bridge cycles (nodes with bridge enabled forming loops) thunderboltBridgeCycles?: string[][]; // Disk usage per node nodeDisk?: Record< string, { total: { inBytes: number }; available: { inBytes: number } } >; } export interface MessageAttachment { type: "image" | "text" | "file" | "generated-image"; name: string; content?: string; preview?: string; mimeType?: string; } export interface TopLogprob { token: string; logprob: number; bytes: number[] | null; } export interface TokenData { token: string; logprob: number; probability: number; topLogprobs: TopLogprob[]; } export interface PrefillProgress { processed: number; total: number; /** Timestamp (performance.now()) when prefill started. */ startedAt: number; } export interface Message { id: string; role: "user" | "assistant" | "system"; content: string; timestamp: number; thinking?: string; attachments?: MessageAttachment[]; ttftMs?: number; // Time to first token in ms (for assistant messages) tps?: number; // Tokens per second (for assistant messages) requestType?: "chat" | "image-generation" | "image-editing"; sourceImageDataUrl?: string; // For image editing regeneration tokens?: TokenData[]; } export interface Conversation { id: string; name: string; messages: Message[]; createdAt: number; updatedAt: number; modelId: string | null; sharding: string | null; instanceType: string | null; enableThinking: boolean | null; } const STORAGE_KEY = "exo-conversations"; const IMAGE_PARAMS_STORAGE_KEY = "exo-image-generation-params"; // Image generation params interface matching backend API export interface ImageGenerationParams { // Basic params size: | "auto" | "512x512" | "768x768" | "1024x1024" | "1024x768" | "768x1024" | "1024x1536" | "1536x1024"; quality: "low" | "medium" | "high"; outputFormat: "png" | "jpeg"; numImages: number; // Streaming params stream: boolean; partialImages: number; // Advanced params seed: number | null; numInferenceSteps: number | null; guidance: number | null; negativePrompt: string | null; numSyncSteps: number | null; // Edit mode params inputFidelity: "low" | "high"; } // Image being edited export interface EditingImage { imageDataUrl: string; sourceMessage: Message; } const DEFAULT_IMAGE_PARAMS: ImageGenerationParams = { size: "auto", quality: "medium", outputFormat: "png", numImages: 1, stream: true, partialImages: 3, seed: null, numInferenceSteps: null, guidance: null, negativePrompt: null, numSyncSteps: null, inputFidelity: "low", }; interface GranularNodeState { nodeIdentities?: Record; nodeMemory?: Record; nodeSystem?: Record; nodeNetwork?: Record; } function transformNetworkInterface(iface: RawNetworkInterfaceInfo): { name?: string; addresses: string[]; } { const addresses: string[] = []; if (iface.ipAddress && typeof iface.ipAddress === "string") { addresses.push(iface.ipAddress); } if (Array.isArray(iface.addresses)) { for (const addr of iface.addresses) { if (typeof addr === "string") addresses.push(addr); else if (addr && typeof addr === "object" && addr.address) addresses.push(addr.address); } } if (Array.isArray(iface.ipAddresses)) { addresses.push( ...iface.ipAddresses.filter((a): a is string => typeof a === "string"), ); } if (Array.isArray(iface.ips)) { addresses.push( ...iface.ips.filter((a): a is string => typeof a === "string"), ); } if (iface.ipv4 && typeof iface.ipv4 === "string") addresses.push(iface.ipv4); if (iface.ipv6 && typeof iface.ipv6 === "string") addresses.push(iface.ipv6); return { name: iface.name, addresses: Array.from(new Set(addresses)), }; } function transformTopology( raw: RawTopology, granularState: GranularNodeState, ): TopologyData { const nodes: Record = {}; const edges: TopologyEdge[] = []; for (const nodeId of raw.nodes || []) { if (!nodeId) continue; // Get data from granular state mappings const identity = granularState.nodeIdentities?.[nodeId]; const memory = granularState.nodeMemory?.[nodeId]; const system = granularState.nodeSystem?.[nodeId]; const network = granularState.nodeNetwork?.[nodeId]; const ramTotal = memory?.ramTotal?.inBytes ?? 0; const ramAvailable = memory?.ramAvailable?.inBytes ?? 0; const ramUsage = Math.max(ramTotal - ramAvailable, 0); const rawInterfaces = network?.interfaces || []; const networkInterfaces = rawInterfaces.map(transformNetworkInterface); const ipToInterface: Record = {}; for (const iface of networkInterfaces) { for (const addr of iface.addresses || []) { ipToInterface[addr] = iface.name ?? ""; } } nodes[nodeId] = { system_info: { model_id: identity?.modelId ?? "Unknown", chip: identity?.chipId, memory: ramTotal, }, network_interfaces: networkInterfaces, ip_to_interface: ipToInterface, macmon_info: { memory: { ram_usage: ramUsage, ram_total: ramTotal, }, temp: system?.temp !== undefined ? { gpu_temp_avg: system.temp } : undefined, gpu_usage: system?.gpuUsage !== undefined ? [0, system.gpuUsage] : undefined, sys_power: system?.sysPower, }, last_macmon_update: Date.now() / 1000, friendly_name: identity?.friendlyName, os_version: identity?.osVersion, }; } // Handle connections - nested mapping format { source: { sink: [edges] } } const connections = raw.connections; if (connections && typeof connections === "object") { for (const [source, sinks] of Object.entries(connections)) { if (!sinks || typeof sinks !== "object") continue; for (const [sink, edgeList] of Object.entries(sinks)) { if (!Array.isArray(edgeList)) continue; for (const edge of edgeList) { let sendBackIp: string | undefined; let sourceRdmaIface: string | undefined; let sinkRdmaIface: string | undefined; if (edge && typeof edge === "object" && "sinkMultiaddr" in edge) { const multiaddr = edge.sinkMultiaddr; if (multiaddr) { sendBackIp = multiaddr.ip_address || extractIpFromMultiaddr(multiaddr.address); } } else if ( edge && typeof edge === "object" && "sourceRdmaIface" in edge ) { sourceRdmaIface = edge.sourceRdmaIface; sinkRdmaIface = edge.sinkRdmaIface; } if (nodes[source] && nodes[sink] && source !== sink) { edges.push({ source, target: sink, sendBackIp, sourceRdmaIface, sinkRdmaIface, }); } } } } } return { nodes, edges }; } function extractIpFromMultiaddr(ma?: string): string | undefined { if (!ma) return undefined; const parts = ma.split("/"); const ip4Idx = parts.indexOf("ip4"); const ip6Idx = parts.indexOf("ip6"); const idx = ip4Idx >= 0 ? ip4Idx : ip6Idx; if (idx >= 0 && parts.length > idx + 1) { return parts[idx + 1]; } return undefined; } class AppStore { // Conversation state conversations = $state([]); activeConversationId = $state(null); // Chat state hasStartedChat = $state(false); messages = $state([]); currentResponse = $state(""); isLoading = $state(false); // Performance metrics ttftMs = $state(null); // Time to first token in ms tps = $state(null); // Tokens per second totalTokens = $state(0); // Total tokens in current response prefillProgress = $state(null); // Abort controller for stopping generation private currentAbortController: AbortController | null = null; // Topology state topologyData = $state(null); instances = $state>({}); runners = $state>({}); downloads = $state>({}); nodeDisk = $state< Record< string, { total: { inBytes: number }; available: { inBytes: number } } > >({}); placementPreviews = $state([]); selectedPreviewModelId = $state(null); isLoadingPreviews = $state(false); previewNodeFilter = $state>(new Set()); lastUpdate = $state(null); nodeIdentities = $state>({}); thunderboltBridgeCycles = $state([]); nodeThunderbolt = $state< Record< string, { interfaces: Array<{ rdmaInterface: string; domainUuid: string; linkSpeed: string; }>; } > >({}); nodeRdmaCtl = $state>({}); nodeThunderboltBridge = $state< Record< string, { enabled: boolean; exists: boolean; serviceName?: string | null } > >({}); // UI state isTopologyMinimized = $state(false); isSidebarOpen = $state(false); // Hidden by default, shown when in chat mode debugMode = $state(false); topologyOnlyMode = $state(false); chatSidebarVisible = $state(true); // Shown by default mobileChatSidebarOpen = $state(false); // Mobile drawer state mobileRightSidebarOpen = $state(false); // Mobile right drawer state // Image generation params imageGenerationParams = $state({ ...DEFAULT_IMAGE_PARAMS, }); // Image editing state editingImage = $state(null); /** True when the backend is reachable. */ isConnected = $state(true); /** Number of consecutive fetch failures. */ private consecutiveFailures = 0; private static readonly CONNECTION_LOST_THRESHOLD = 3; private fetchInterval: ReturnType | null = null; private previewsInterval: ReturnType | null = null; private lastConversationPersistTs = 0; private previousNodeIds: Set = new Set(); constructor() { if (browser) { this.startPolling(); this.loadConversationsFromStorage(); this.loadDebugModeFromStorage(); this.loadTopologyOnlyModeFromStorage(); this.loadChatSidebarVisibleFromStorage(); this.loadImageGenerationParamsFromStorage(); } } /** * Load conversations from localStorage */ private loadConversationsFromStorage() { try { const stored = localStorage.getItem(STORAGE_KEY); if (stored) { const parsed = JSON.parse(stored) as Array>; this.conversations = parsed.map((conversation) => ({ id: conversation.id ?? generateUUID(), name: conversation.name ?? "Chat", messages: conversation.messages ?? [], createdAt: conversation.createdAt ?? Date.now(), updatedAt: conversation.updatedAt ?? Date.now(), modelId: conversation.modelId ?? null, sharding: conversation.sharding ?? null, instanceType: conversation.instanceType ?? null, enableThinking: conversation.enableThinking ?? null, })); } } catch (error) { console.error("Failed to load conversations:", error); } } /** * Save conversations to localStorage */ private saveConversationsToStorage() { try { // Strip tokens from messages before saving to avoid bloating localStorage const stripped = this.conversations.map((conv) => ({ ...conv, messages: conv.messages.map((msg) => { if (msg.tokens) { const { tokens: _, ...rest } = msg; return rest; } return msg; }), })); localStorage.setItem(STORAGE_KEY, JSON.stringify(stripped)); } catch (error) { console.error("Failed to save conversations:", error); } } private loadDebugModeFromStorage() { try { const stored = localStorage.getItem("exo-debug-mode"); if (stored !== null) { this.debugMode = stored === "true"; } } catch (error) { console.error("Failed to load debug mode:", error); } } private saveDebugModeToStorage() { try { localStorage.setItem("exo-debug-mode", this.debugMode ? "true" : "false"); } catch (error) { console.error("Failed to save debug mode:", error); } } private loadTopologyOnlyModeFromStorage() { try { const stored = localStorage.getItem("exo-topology-only-mode"); if (stored !== null) { this.topologyOnlyMode = stored === "true"; } } catch (error) { console.error("Failed to load topology only mode:", error); } } private saveTopologyOnlyModeToStorage() { try { localStorage.setItem( "exo-topology-only-mode", this.topologyOnlyMode ? "true" : "false", ); } catch (error) { console.error("Failed to save topology only mode:", error); } } private loadChatSidebarVisibleFromStorage() { try { const stored = localStorage.getItem("exo-chat-sidebar-visible"); if (stored !== null) { this.chatSidebarVisible = stored === "true"; } } catch (error) { console.error("Failed to load chat sidebar visibility:", error); } } private saveChatSidebarVisibleToStorage() { try { localStorage.setItem( "exo-chat-sidebar-visible", this.chatSidebarVisible ? "true" : "false", ); } catch (error) { console.error("Failed to save chat sidebar visibility:", error); } } private loadImageGenerationParamsFromStorage() { try { const stored = localStorage.getItem(IMAGE_PARAMS_STORAGE_KEY); if (stored) { const parsed = JSON.parse(stored) as Partial; this.imageGenerationParams = { ...DEFAULT_IMAGE_PARAMS, ...parsed, }; } } catch (error) { console.error("Failed to load image generation params:", error); } } private saveImageGenerationParamsToStorage() { try { localStorage.setItem( IMAGE_PARAMS_STORAGE_KEY, JSON.stringify(this.imageGenerationParams), ); } catch (error) { console.error("Failed to save image generation params:", error); } } getImageGenerationParams(): ImageGenerationParams { return this.imageGenerationParams; } setImageGenerationParams(params: Partial) { this.imageGenerationParams = { ...this.imageGenerationParams, ...params, }; this.saveImageGenerationParamsToStorage(); } resetImageGenerationParams() { this.imageGenerationParams = { ...DEFAULT_IMAGE_PARAMS }; this.saveImageGenerationParamsToStorage(); } setEditingImage(imageDataUrl: string, sourceMessage: Message) { this.editingImage = { imageDataUrl, sourceMessage }; } clearEditingImage() { this.editingImage = null; } /** * Create a new conversation */ createConversation(name?: string): string { const id = generateUUID(); const now = Date.now(); // Try to derive model and strategy immediately from selected model or running instances let derivedModelId = this.selectedChatModel || null; let derivedInstanceType: string | null = null; let derivedSharding: string | null = null; // If no selected model, fall back to the first running instance if (!derivedModelId) { const firstInstance = Object.values(this.instances)[0]; if (firstInstance) { const candidateModel = this.extractInstanceModelId(firstInstance); derivedModelId = candidateModel ?? null; const details = this.describeInstance(firstInstance); derivedInstanceType = details.instanceType; derivedSharding = details.sharding; } } else { // If selected model is set, attempt to get its details from instances for (const [, instanceWrapper] of Object.entries(this.instances)) { const candidateModelId = this.extractInstanceModelId(instanceWrapper); if (candidateModelId === derivedModelId) { const details = this.describeInstance(instanceWrapper); derivedInstanceType = details.instanceType; derivedSharding = details.sharding; break; } } } const conversation: Conversation = { id, name: name || `Chat ${new Date(now).toLocaleString("en-US", { month: "short", day: "numeric", hour: "2-digit", minute: "2-digit" })}`, messages: [], createdAt: now, updatedAt: now, modelId: derivedModelId, sharding: derivedSharding, instanceType: derivedInstanceType, enableThinking: null, }; this.conversations.unshift(conversation); this.activeConversationId = id; this.messages = []; this.hasStartedChat = true; this.isTopologyMinimized = true; this.isSidebarOpen = true; // Auto-open sidebar when chatting this.saveConversationsToStorage(); return id; } /** * Load a conversation by ID */ loadConversation(id: string): boolean { const conversation = this.conversations.find((c) => c.id === id); if (!conversation) return false; this.activeConversationId = id; this.messages = [...conversation.messages]; this.hasStartedChat = true; this.isTopologyMinimized = true; this.isSidebarOpen = true; // Auto-open sidebar when chatting this.thinkingEnabled = conversation.enableThinking ?? true; this.refreshConversationModelFromInstances(); // Sync global selection to the loaded conversation's model so reactive // effects in +page.svelte can determine the correct chat launch state. this.selectedChatModel = conversation.modelId || ""; return true; } /** * Delete a conversation by ID */ deleteConversation(id: string) { this.conversations = this.conversations.filter((c) => c.id !== id); if (this.activeConversationId === id) { this.activeConversationId = null; this.messages = []; this.hasStartedChat = false; this.isTopologyMinimized = false; } this.saveConversationsToStorage(); } /** * Delete all conversations */ deleteAllConversations() { this.conversations = []; this.activeConversationId = null; this.messages = []; this.hasStartedChat = false; this.isTopologyMinimized = false; this.saveConversationsToStorage(); } /** * Rename a conversation */ renameConversation(id: string, newName: string) { const conversation = this.conversations.find((c) => c.id === id); if (conversation) { conversation.name = newName; conversation.updatedAt = Date.now(); this.saveConversationsToStorage(); } } private getTaggedValue(obj: unknown): [string | null, unknown] { if (!obj || typeof obj !== "object") return [null, null]; const keys = Object.keys(obj as Record); if (keys.length === 1) { return [keys[0], (obj as Record)[keys[0]]]; } return [null, null]; } private extractInstanceModelId(instanceWrapped: unknown): string | null { const [, instance] = this.getTaggedValue(instanceWrapped); if (!instance || typeof instance !== "object") return null; const inst = instance as { shardAssignments?: { modelId?: string } }; return inst.shardAssignments?.modelId ?? null; } private describeInstance(instanceWrapped: unknown): { sharding: string | null; instanceType: string | null; } { const [instanceTag, instance] = this.getTaggedValue(instanceWrapped); if (!instance || typeof instance !== "object") { return { sharding: null, instanceType: null }; } let instanceType: string | null = null; if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring"; else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA"; let sharding: string | null = null; const inst = instance as { shardAssignments?: { runnerToShard?: Record }; }; const runnerToShard = inst.shardAssignments?.runnerToShard || {}; const firstShardWrapped = Object.values(runnerToShard)[0]; if (firstShardWrapped) { const [shardTag] = this.getTaggedValue(firstShardWrapped); if (shardTag === "PipelineShardMetadata") sharding = "Pipeline"; else if (shardTag === "TensorShardMetadata") sharding = "Tensor"; else if (shardTag === "PrefillDecodeShardMetadata") sharding = "Prefill/Decode"; } return { sharding, instanceType }; } private buildConversationModelInfo(modelId: string): { modelId: string; sharding: string | null; instanceType: string | null; } { let sharding: string | null = null; let instanceType: string | null = null; for (const [, instanceWrapper] of Object.entries(this.instances)) { const candidateModelId = this.extractInstanceModelId(instanceWrapper); if (candidateModelId === modelId) { const details = this.describeInstance(instanceWrapper); sharding = details.sharding; instanceType = details.instanceType; break; } } return { modelId, sharding, instanceType }; } private applyConversationModelInfo(info: { modelId: string; sharding: string | null; instanceType: string | null; }) { if (!this.activeConversationId) return; const conversation = this.conversations.find( (c) => c.id === this.activeConversationId, ); if (!conversation) return; // Keep the first known modelId stable; only backfill if missing if (!conversation.modelId) { conversation.modelId = info.modelId; } conversation.sharding = info.sharding; conversation.instanceType = info.instanceType; this.saveConversationsToStorage(); } private getModelTail(modelId: string): string { const parts = modelId.split("/"); return (parts[parts.length - 1] || modelId).toLowerCase(); } private isBetterModelId( currentId: string | null, candidateId: string | null, ): boolean { if (!candidateId) return false; if (!currentId) return true; const currentTail = this.getModelTail(currentId); const candidateTail = this.getModelTail(candidateId); return ( candidateTail.length > currentTail.length && candidateTail.startsWith(currentTail) ); } private refreshConversationModelFromInstances() { if (!this.activeConversationId) return; const conversation = this.conversations.find( (c) => c.id === this.activeConversationId, ); if (!conversation) return; // Prefer stored model; do not replace it once set. Only backfill when missing. let modelId = conversation.modelId; // If missing, try the selected model if (!modelId && this.selectedChatModel) { modelId = this.selectedChatModel; } // If still missing, fall back to first instance model if (!modelId) { const firstInstance = Object.values(this.instances)[0]; if (firstInstance) { modelId = this.extractInstanceModelId(firstInstance); } } if (!modelId) return; // If a more specific instance modelId is available (e.g., adds "-4bit"), prefer it let preferredModelId = modelId; for (const [, instanceWrapper] of Object.entries(this.instances)) { const candidate = this.extractInstanceModelId(instanceWrapper); if (!candidate) continue; if (candidate === preferredModelId) { break; } if (this.isBetterModelId(preferredModelId, candidate)) { preferredModelId = candidate; } } if (this.isBetterModelId(conversation.modelId, preferredModelId)) { conversation.modelId = preferredModelId; } const info = this.buildConversationModelInfo(preferredModelId); const hasNewInfo = Boolean( info.sharding || info.instanceType || !conversation.modelId, ); if (hasNewInfo) { this.applyConversationModelInfo(info); } } getDebugMode(): boolean { return this.debugMode; } /** * Update the active conversation with current messages */ private updateActiveConversation() { if (!this.activeConversationId) return; const conversation = this.conversations.find( (c) => c.id === this.activeConversationId, ); if (conversation) { conversation.messages = [...this.messages]; conversation.updatedAt = Date.now(); // Auto-generate name from first user message if still has default name if (conversation.name.startsWith("Chat ")) { const firstUserMsg = conversation.messages.find( (m) => m.role === "user" && m.content.trim(), ); if (firstUserMsg) { // Clean up the content - remove file context markers and whitespace let content = firstUserMsg.content .replace(/\[File:.*?\][\s\S]*?```[\s\S]*?```/g, "") // Remove file attachments .trim(); if (content) { const preview = content.slice(0, 50); conversation.name = preview.length < content.length ? preview + "..." : preview; } } } this.saveConversationsToStorage(); } } private persistActiveConversation(throttleMs = 400) { const now = Date.now(); if (now - this.lastConversationPersistTs < throttleMs) return; this.lastConversationPersistTs = now; this.updateActiveConversation(); } /** * Update a message in a specific conversation by ID. * Returns false if conversation or message not found. */ private updateConversationMessage( conversationId: string, messageId: string, updater: (message: Message) => void, ): boolean { const conversation = this.conversations.find( (c) => c.id === conversationId, ); if (!conversation) return false; const message = conversation.messages.find((m) => m.id === messageId); if (!message) return false; updater(message); return true; } /** * Sync this.messages from the target conversation if it matches the active conversation. */ private syncActiveMessagesIfNeeded(conversationId: string): void { if (this.activeConversationId === conversationId) { const conversation = this.conversations.find( (c) => c.id === conversationId, ); if (conversation) { this.messages = [...conversation.messages]; } } } /** * Check if a conversation still exists. */ private conversationExists(conversationId: string): boolean { return this.conversations.some((c) => c.id === conversationId); } /** * Persist a specific conversation to storage. */ private persistConversation(conversationId: string, throttleMs = 400): void { const now = Date.now(); if (now - this.lastConversationPersistTs < throttleMs) return; this.lastConversationPersistTs = now; const conversation = this.conversations.find( (c) => c.id === conversationId, ); if (conversation) { conversation.updatedAt = Date.now(); // Auto-generate name from first user message if still has default name if (conversation.name.startsWith("Chat ")) { const firstUserMsg = conversation.messages.find( (m) => m.role === "user" && m.content.trim(), ); if (firstUserMsg) { let content = firstUserMsg.content .replace(/\[File:.*?\][\s\S]*?```[\s\S]*?```/g, "") .trim(); if (content) { const preview = content.slice(0, 50); conversation.name = preview.length < content.length ? preview + "..." : preview; } } } this.saveConversationsToStorage(); } } /** * Add a message directly to a specific conversation. * Returns the message if added, null if conversation not found. */ private addMessageToConversation( conversationId: string, role: "user" | "assistant", content: string, ): Message | null { const conversation = this.conversations.find( (c) => c.id === conversationId, ); if (!conversation) return null; const message: Message = { id: generateUUID(), role, content, timestamp: Date.now(), }; conversation.messages.push(message); return message; } /** * Toggle sidebar visibility */ toggleSidebar() { this.isSidebarOpen = !this.isSidebarOpen; } setDebugMode(enabled: boolean) { this.debugMode = enabled; this.saveDebugModeToStorage(); } toggleDebugMode() { this.debugMode = !this.debugMode; this.saveDebugModeToStorage(); } getTopologyOnlyMode(): boolean { return this.topologyOnlyMode; } setTopologyOnlyMode(enabled: boolean) { this.topologyOnlyMode = enabled; this.saveTopologyOnlyModeToStorage(); } toggleTopologyOnlyMode() { this.topologyOnlyMode = !this.topologyOnlyMode; this.saveTopologyOnlyModeToStorage(); } getChatSidebarVisible(): boolean { return this.chatSidebarVisible; } setChatSidebarVisible(visible: boolean) { this.chatSidebarVisible = visible; this.saveChatSidebarVisibleToStorage(); } toggleChatSidebarVisible() { this.chatSidebarVisible = !this.chatSidebarVisible; this.saveChatSidebarVisibleToStorage(); } getMobileChatSidebarOpen(): boolean { return this.mobileChatSidebarOpen; } setMobileChatSidebarOpen(open: boolean) { this.mobileChatSidebarOpen = open; } toggleMobileChatSidebar() { this.mobileChatSidebarOpen = !this.mobileChatSidebarOpen; } getMobileRightSidebarOpen(): boolean { return this.mobileRightSidebarOpen; } setMobileRightSidebarOpen(open: boolean) { this.mobileRightSidebarOpen = open; } toggleMobileRightSidebar() { this.mobileRightSidebarOpen = !this.mobileRightSidebarOpen; } startPolling() { this.fetchState(); this.fetchInterval = setInterval(() => this.fetchState(), 1000); } stopPolling() { if (this.fetchInterval) { clearInterval(this.fetchInterval); this.fetchInterval = null; } this.stopPreviewsPolling(); } async fetchState() { try { const response = await fetch("/state"); if (!response.ok) { throw new Error(`Failed to fetch state: ${response.status}`); } const data: RawStateResponse = await response.json(); if (data.topology) { this.topologyData = transformTopology(data.topology, { nodeIdentities: data.nodeIdentities, nodeMemory: data.nodeMemory, nodeSystem: data.nodeSystem, nodeNetwork: data.nodeNetwork, }); // Handle topology changes for preview filter this.handleTopologyChange(); } if (data.instances) { this.instances = data.instances; this.refreshConversationModelFromInstances(); } if (data.runners) { this.runners = data.runners; } if (data.downloads) { this.downloads = data.downloads; } if (data.nodeDisk) { this.nodeDisk = data.nodeDisk; } // Node identities (for OS version mismatch detection) this.nodeIdentities = data.nodeIdentities ?? {}; // Thunderbolt identifiers per node this.nodeThunderbolt = data.nodeThunderbolt ?? {}; // RDMA ctl status per node this.nodeRdmaCtl = data.nodeRdmaCtl ?? {}; // Thunderbolt bridge cycles this.thunderboltBridgeCycles = data.thunderboltBridgeCycles ?? []; // Thunderbolt bridge status per node this.nodeThunderboltBridge = data.nodeThunderboltBridge ?? {}; this.lastUpdate = Date.now(); // Connection recovered if (!this.isConnected) { this.isConnected = true; } this.consecutiveFailures = 0; } catch (error) { this.consecutiveFailures++; if ( this.consecutiveFailures >= AppStore.CONNECTION_LOST_THRESHOLD && this.isConnected ) { this.isConnected = false; } console.error("Error fetching state:", error); } } async fetchPlacementPreviews(modelId: string, showLoading = true) { if (!modelId) return; if (showLoading) { this.isLoadingPreviews = true; } this.selectedPreviewModelId = modelId; try { let url = `/instance/previews?model_id=${encodeURIComponent(modelId)}`; // Add node filter if active if (this.previewNodeFilter.size > 0) { for (const nodeId of this.previewNodeFilter) { url += `&node_ids=${encodeURIComponent(nodeId)}`; } } const response = await fetch(url); if (!response.ok) { throw new Error( `Failed to fetch placement previews: ${response.status}`, ); } const data: PlacementPreviewResponse = await response.json(); this.placementPreviews = data.previews; } catch (error) { console.error("Error fetching placement previews:", error); this.placementPreviews = []; } finally { if (showLoading) { this.isLoadingPreviews = false; } } } startPreviewsPolling(modelId: string) { // Stop any existing preview polling this.stopPreviewsPolling(); // Fetch immediately this.fetchPlacementPreviews(modelId); // Then poll every 15 seconds (don't show loading spinner for subsequent fetches) this.previewsInterval = setInterval(() => { if (this.selectedPreviewModelId) { this.fetchPlacementPreviews(this.selectedPreviewModelId, false); } }, 15000); } stopPreviewsPolling() { if (this.previewsInterval) { clearInterval(this.previewsInterval); this.previewsInterval = null; } } selectPreviewModel(modelId: string | null) { if (modelId) { this.startPreviewsPolling(modelId); } else { this.stopPreviewsPolling(); this.selectedPreviewModelId = null; this.placementPreviews = []; } } /** * Toggle a node in the preview filter and re-fetch placements */ togglePreviewNodeFilter(nodeId: string) { const next = new Set(this.previewNodeFilter); if (next.has(nodeId)) { next.delete(nodeId); } else { next.add(nodeId); } this.previewNodeFilter = next; // Re-fetch with new filter if we have a selected model if (this.selectedPreviewModelId) { this.fetchPlacementPreviews(this.selectedPreviewModelId, false); } } /** * Clear the preview node filter and re-fetch placements */ clearPreviewNodeFilter() { this.previewNodeFilter = new Set(); // Re-fetch with no filter if we have a selected model if (this.selectedPreviewModelId) { this.fetchPlacementPreviews(this.selectedPreviewModelId, false); } } /** * Handle topology changes - clean up filter and re-fetch if needed */ private handleTopologyChange() { if (!this.topologyData) return; const currentNodeIds = new Set(Object.keys(this.topologyData.nodes)); // Check if nodes have changed const nodesAdded = [...currentNodeIds].some( (id) => !this.previousNodeIds.has(id), ); const nodesRemoved = [...this.previousNodeIds].some( (id) => !currentNodeIds.has(id), ); if (nodesAdded || nodesRemoved) { // Clean up filter - remove any nodes that no longer exist if (this.previewNodeFilter.size > 0) { const validFilterNodes = new Set( [...this.previewNodeFilter].filter((id) => currentNodeIds.has(id)), ); if (validFilterNodes.size !== this.previewNodeFilter.size) { this.previewNodeFilter = validFilterNodes; } } // Re-fetch previews if we have a selected model (topology changed) if (this.selectedPreviewModelId) { this.fetchPlacementPreviews(this.selectedPreviewModelId, false); } } // Update tracked node IDs for next comparison this.previousNodeIds = currentNodeIds; } /** * Starts a chat conversation - triggers the topology minimization animation * Creates a new conversation if none is active */ startChat() { if (!this.activeConversationId) { this.createConversation(); } else { this.hasStartedChat = true; this.isSidebarOpen = true; // Auto-open sidebar when chatting // Small delay before minimizing for a nice visual effect setTimeout(() => { this.isTopologyMinimized = true; }, 100); } } /** * Add a message to the conversation */ addMessage(role: "user" | "assistant", content: string) { const message: Message = { id: generateUUID(), role, content, timestamp: Date.now(), }; this.messages.push(message); return message; } /** * Delete a message and all subsequent messages */ deleteMessage(messageId: string) { const messageIndex = this.messages.findIndex((m) => m.id === messageId); if (messageIndex === -1) return; // Remove this message and all subsequent messages this.messages = this.messages.slice(0, messageIndex); this.updateActiveConversation(); } /** * Edit a user message content (does not regenerate response) */ editMessage(messageId: string, newContent: string) { const message = this.messages.find((m) => m.id === messageId); if (!message) return; message.content = newContent; message.timestamp = Date.now(); this.updateActiveConversation(); } /** * Edit a user message and regenerate the response */ async editAndRegenerate( messageId: string, newContent: string, ): Promise { const messageIndex = this.messages.findIndex((m) => m.id === messageId); if (messageIndex === -1) return; const message = this.messages[messageIndex]; if (message.role !== "user") return; // Update the message content message.content = newContent; message.timestamp = Date.now(); // Remove all messages after this one (including the assistant response) this.messages = this.messages.slice(0, messageIndex + 1); // Regenerate the response await this.regenerateLastResponse(); } /** * Regenerate the last assistant response */ async regenerateLastResponse(): Promise { if (this.isLoading) return; // Find the last user message let lastUserIndex = -1; for (let i = this.messages.length - 1; i >= 0; i--) { if (this.messages[i].role === "user") { lastUserIndex = i; break; } } if (lastUserIndex === -1) return; const lastUserMessage = this.messages[lastUserIndex]; const requestType = lastUserMessage.requestType || "chat"; const prompt = lastUserMessage.content; // Remove messages after user message (including the user message for image requests // since generateImage/editImage will re-add it) this.messages = this.messages.slice(0, lastUserIndex); this.updateActiveConversation(); switch (requestType) { case "image-generation": await this.generateImage(prompt); break; case "image-editing": if (lastUserMessage.sourceImageDataUrl) { await this.editImage(prompt, lastUserMessage.sourceImageDataUrl); } else { // Can't regenerate edit without source image - restore user message and show error this.messages.push(lastUserMessage); const errorMessage = this.addMessage("assistant", ""); const idx = this.messages.findIndex((m) => m.id === errorMessage.id); if (idx !== -1) { this.messages[idx].content = "Error: Cannot regenerate image edit - source image not found"; } this.updateActiveConversation(); } break; case "chat": default: // Restore the user message for chat regeneration this.messages.push(lastUserMessage); await this.regenerateChatCompletion(); break; } } /** * Regenerate response from a specific token index. * Truncates the assistant message at the given token and re-generates from there. */ async regenerateFromToken( messageId: string, tokenIndex: number, ): Promise { if (this.isLoading) return; const targetConversationId = this.activeConversationId; if (!targetConversationId) return; const msgIndex = this.messages.findIndex((m) => m.id === messageId); if (msgIndex === -1) return; const msg = this.messages[msgIndex]; if ( msg.role !== "assistant" || !msg.tokens || tokenIndex >= msg.tokens.length ) return; // Keep tokens up to (not including) the specified index const tokensToKeep = msg.tokens.slice(0, tokenIndex); const prefixText = tokensToKeep.map((t) => t.token).join(""); // Remove all messages after this assistant message this.messages = this.messages.slice(0, msgIndex + 1); // Update the message to show the prefix this.messages[msgIndex].content = prefixText; this.messages[msgIndex].tokens = tokensToKeep; this.updateActiveConversation(); // Set up for continuation - modify the existing message in place this.isLoading = true; this.currentResponse = prefixText; this.ttftMs = null; this.tps = null; this.totalTokens = tokensToKeep.length; try { // Build messages for API - include the partial assistant message const systemPrompt = { role: "system" as const, content: "You are a helpful AI assistant. Respond directly and concisely. Do not show your reasoning or thought process.", }; const apiMessages = [ systemPrompt, ...this.messages.map((m) => { let msgContent = m.content; if (m.attachments) { for (const attachment of m.attachments) { if (attachment.type === "text" && attachment.content) { msgContent += `\n\n[File: ${attachment.name}]\n\`\`\`\n${attachment.content}\n\`\`\``; } } } return { role: m.role, content: msgContent }; }), ]; const modelToUse = this.getModelForRequest(); if (!modelToUse) { throw new Error("No model available"); } const requestStartTime = performance.now(); let firstTokenTime: number | null = null; let tokenCount = tokensToKeep.length; const response = await fetch("/v1/chat/completions", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ model: modelToUse, messages: apiMessages, stream: true, logprobs: true, top_logprobs: 5, }), }); if (!response.ok) { const errorText = await response.text(); throw new Error(`API error: ${response.status} - ${errorText}`); } const reader = response.body?.getReader(); if (!reader) throw new Error("No response body"); let fullContent = prefixText; let streamedThinking = ""; const collectedTokens: TokenData[] = [...tokensToKeep]; interface ChatCompletionChunk { choices?: Array<{ delta?: { content?: string; reasoning_content?: string }; logprobs?: { content?: Array<{ token: string; logprob: number; top_logprobs?: Array<{ token: string; logprob: number; bytes: number[] | null; }>; }>; }; }>; } await this.parseSSEStream( reader, targetConversationId, (parsed) => { const choice = parsed.choices?.[0]; const delta = choice?.delta?.content; const thinkingDelta = choice?.delta?.reasoning_content; // Collect logprobs data const logprobsContent = choice?.logprobs?.content; if (logprobsContent) { for (const item of logprobsContent) { collectedTokens.push({ token: item.token, logprob: item.logprob, probability: Math.exp(item.logprob), topLogprobs: (item.top_logprobs || []).map((t) => ({ token: t.token, logprob: t.logprob, bytes: t.bytes, })), }); } } if (thinkingDelta) { streamedThinking += thinkingDelta; } if (delta || thinkingDelta) { if (firstTokenTime === null) { firstTokenTime = performance.now(); this.ttftMs = firstTokenTime - requestStartTime; } tokenCount += 1; this.totalTokens = tokenCount; if (firstTokenTime !== null && tokenCount > tokensToKeep.length) { const elapsed = performance.now() - firstTokenTime; this.tps = ((tokenCount - tokensToKeep.length) / elapsed) * 1000; } if (delta) { fullContent += delta; } const { displayContent, thinkingContent: tagThinking } = this.stripThinkingTags(fullContent); const combinedThinking = [streamedThinking, tagThinking] .filter(Boolean) .join("\n\n"); if (this.activeConversationId === targetConversationId) { this.currentResponse = displayContent; } // Update existing message in place this.updateConversationMessage( targetConversationId, messageId, (m) => { m.content = displayContent; m.thinking = combinedThinking || undefined; m.tokens = [...collectedTokens]; }, ); this.syncActiveMessagesIfNeeded(targetConversationId); this.persistConversation(targetConversationId); } }, ); // Final update if (this.conversationExists(targetConversationId)) { const { displayContent, thinkingContent: tagThinking } = this.stripThinkingTags(fullContent); const finalThinking = [streamedThinking, tagThinking] .filter(Boolean) .join("\n\n"); this.updateConversationMessage(targetConversationId, messageId, (m) => { m.content = displayContent; m.thinking = finalThinking || undefined; m.tokens = [...collectedTokens]; if (this.ttftMs !== null) m.ttftMs = this.ttftMs; if (this.tps !== null) m.tps = this.tps; }); this.syncActiveMessagesIfNeeded(targetConversationId); this.persistConversation(targetConversationId); } } catch (error) { console.error("Error regenerating from token:", error); if (this.conversationExists(targetConversationId)) { this.updateConversationMessage(targetConversationId, messageId, (m) => { m.content = `${prefixText}\n\nError: ${error instanceof Error ? error.message : "Unknown error"}`; }); this.syncActiveMessagesIfNeeded(targetConversationId); this.persistConversation(targetConversationId); } } finally { this.isLoading = false; this.currentResponse = ""; this.saveConversationsToStorage(); } } /** * Helper method to regenerate a chat completion response */ private async regenerateChatCompletion(): Promise { // Capture the target conversation ID at the start of the request const targetConversationId = this.activeConversationId; if (!targetConversationId) return; const targetConversation = this.conversations.find( (c) => c.id === targetConversationId, ); if (!targetConversation) return; this.isLoading = true; this.currentResponse = ""; // Create placeholder for assistant message directly in target conversation const assistantMessage = this.addMessageToConversation( targetConversationId, "assistant", "", ); if (!assistantMessage) { this.isLoading = false; return; } // Sync to this.messages if viewing the target conversation this.syncActiveMessagesIfNeeded(targetConversationId); try { const systemPrompt = { role: "system" as const, content: "You are a helpful AI assistant. Respond directly and concisely. Do not show your reasoning or thought process.", }; const apiMessages = [ systemPrompt, ...targetConversation.messages.slice(0, -1).map((m) => { return { role: m.role, content: m.content }; }), ]; // Determine which model to use const modelToUse = this.getModelForRequest(); if (!modelToUse) { this.updateConversationMessage( targetConversationId, assistantMessage.id, (msg) => { msg.content = "No model is loaded yet. Select a model from the sidebar to get started — it will download and load automatically."; }, ); this.syncActiveMessagesIfNeeded(targetConversationId); this.isLoading = false; this.saveConversationsToStorage(); return; } const response = await fetch("/v1/chat/completions", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ model: modelToUse, messages: apiMessages, stream: true, logprobs: true, top_logprobs: 5, }), }); if (!response.ok) { const errorText = await response.text(); throw new Error(`${response.status} - ${errorText}`); } const reader = response.body?.getReader(); if (!reader) { throw new Error("No response stream available"); } let streamedContent = ""; let streamedThinking = ""; const collectedTokens: TokenData[] = []; interface ChatCompletionChunk { choices?: Array<{ delta?: { content?: string; reasoning_content?: string }; logprobs?: { content?: Array<{ token: string; logprob: number; top_logprobs?: Array<{ token: string; logprob: number; bytes: number[] | null; }>; }>; }; }>; } await this.parseSSEStream( reader, targetConversationId, (parsed) => { const choice = parsed.choices?.[0]; const delta = choice?.delta?.content; const thinkingDelta = choice?.delta?.reasoning_content; // Collect logprobs data const logprobsContent = choice?.logprobs?.content; if (logprobsContent) { for (const item of logprobsContent) { collectedTokens.push({ token: item.token, logprob: item.logprob, probability: Math.exp(item.logprob), topLogprobs: (item.top_logprobs || []).map((t) => ({ token: t.token, logprob: t.logprob, bytes: t.bytes, })), }); } } if (thinkingDelta) { streamedThinking += thinkingDelta; } if (delta || thinkingDelta) { if (delta) { streamedContent += delta; } const { displayContent, thinkingContent: tagThinking } = this.stripThinkingTags(streamedContent); const combinedThinking = [streamedThinking, tagThinking] .filter(Boolean) .join("\n\n"); // Only update currentResponse if target conversation is active if (this.activeConversationId === targetConversationId) { this.currentResponse = displayContent; } // Update the assistant message in the target conversation this.updateConversationMessage( targetConversationId, assistantMessage.id, (msg) => { msg.content = displayContent; msg.thinking = combinedThinking || undefined; msg.tokens = [...collectedTokens]; }, ); this.syncActiveMessagesIfNeeded(targetConversationId); this.persistConversation(targetConversationId); } }, ); // Final cleanup of the message (if conversation still exists) if (this.conversationExists(targetConversationId)) { const { displayContent, thinkingContent: tagThinking } = this.stripThinkingTags(streamedContent); const finalThinking = [streamedThinking, tagThinking] .filter(Boolean) .join("\n\n"); this.updateConversationMessage( targetConversationId, assistantMessage.id, (msg) => { msg.content = displayContent; msg.thinking = finalThinking || undefined; msg.tokens = [...collectedTokens]; }, ); this.syncActiveMessagesIfNeeded(targetConversationId); this.persistConversation(targetConversationId); } } catch (error) { this.handleStreamingError( error, targetConversationId, assistantMessage.id, "Unknown error", ); } finally { this.isLoading = false; this.currentResponse = ""; this.saveConversationsToStorage(); } } /** * Whether thinking is enabled for the current conversation */ thinkingEnabled = $state(true); /** * Selected model for chat (can be set by the UI) */ selectedChatModel = $state(""); /** * Set the model to use for chat */ setSelectedModel(modelId: string) { this.selectedChatModel = modelId; // Clear stats when model changes this.ttftMs = null; this.tps = null; } /** * Strip thinking tags from content for display. * Handles both complete ... blocks and in-progress ... blocks during streaming. */ private stripThinkingTags(content: string): { displayContent: string; thinkingContent: string; } { const extracted: string[] = []; let displayContent = content; // Extract complete ... blocks const completeBlockRegex = /([\s\S]*?)<\/think>/gi; let match: RegExpExecArray | null; while ((match = completeBlockRegex.exec(content)) !== null) { const inner = match[1]?.trim(); if (inner) extracted.push(inner); } displayContent = displayContent.replace(completeBlockRegex, ""); // Handle in-progress thinking block (has but no closing yet) const openTagIndex = displayContent.lastIndexOf(""); if (openTagIndex !== -1) { const inProgressThinking = displayContent.slice(openTagIndex + 7).trim(); if (inProgressThinking) { extracted.push(inProgressThinking); } displayContent = displayContent.slice(0, openTagIndex); } return { displayContent: displayContent.trim(), thinkingContent: extracted.join("\n\n"), }; } /** * Parse an SSE stream and invoke a callback for each parsed JSON chunk. * Handles buffering, line splitting, and conversation deletion checks. * * @param reader - The stream reader from fetch response.body.getReader() * @param targetConversationId - The conversation ID to check for deletion * @param onChunk - Callback invoked with each parsed JSON object from the stream */ private async parseSSEStream( reader: ReadableStreamDefaultReader, targetConversationId: string, onChunk: (parsed: T) => void, onEvent?: Record void>, ): Promise { const decoder = new TextDecoder(); let buffer = ""; while (true) { const { done, value } = await reader.read(); if (done) break; if (!this.conversationExists(targetConversationId)) { break; } buffer += decoder.decode(value, { stream: true }); const lines = buffer.split("\n"); buffer = lines.pop() || ""; for (const line of lines) { const trimmed = line.trim(); if (!trimmed) continue; // Handle SSE comments (": key json") for prefill progress etc. if (trimmed.startsWith(": ") && onEvent) { const comment = trimmed.slice(2); const spaceIdx = comment.indexOf(" "); if (spaceIdx > 0) { const key = comment.slice(0, spaceIdx); if (onEvent[key]) { try { const parsed = JSON.parse(comment.slice(spaceIdx + 1)); onEvent[key](parsed); } catch { // Skip malformed JSON in comment } } } continue; } if (trimmed.startsWith("data: ")) { const data = trimmed.slice(6); if (data === "[DONE]") continue; try { const parsed = JSON.parse(data) as T; onChunk(parsed); } catch { // Skip malformed JSON } } } } // Process any remaining data in the buffer if (buffer.trim() && this.conversationExists(targetConversationId)) { const trimmed = buffer.trim(); if (trimmed.startsWith("data: ") && trimmed.slice(6) !== "[DONE]") { try { const parsed = JSON.parse(trimmed.slice(6)) as T; onChunk(parsed); } catch { // Skip malformed JSON } } } } /** * Handle streaming errors by updating the assistant message with an error. * * @param error - The caught error * @param targetConversationId - The conversation ID * @param assistantMessageId - The assistant message ID to update * @param errorPrefix - Optional prefix for the error message (e.g., "Failed to generate image") */ private handleStreamingError( error: unknown, targetConversationId: string, assistantMessageId: string, errorPrefix = "Failed to get response", ): void { if (this.conversationExists(targetConversationId)) { this.updateConversationMessage( targetConversationId, assistantMessageId, (msg) => { msg.content = `Error: ${error instanceof Error ? error.message : errorPrefix}`; }, ); this.syncActiveMessagesIfNeeded(targetConversationId); this.persistConversation(targetConversationId); } } /** * Get the model to use for a request. * Prefers the provided modelId, then selectedChatModel, then falls back to the first running instance. * * @param modelId - Optional explicit model ID * @returns The model ID to use, or null if none available */ private getModelForRequest(modelId?: string): string | null { if (modelId) return modelId; if (this.selectedChatModel) return this.selectedChatModel; // Try to get model from first running instance for (const [, instanceWrapper] of Object.entries(this.instances)) { if (instanceWrapper && typeof instanceWrapper === "object") { const keys = Object.keys(instanceWrapper as Record); if (keys.length === 1) { const instance = (instanceWrapper as Record)[ keys[0] ] as { shardAssignments?: { modelId?: string } }; if (instance?.shardAssignments?.modelId) { return instance.shardAssignments.modelId; } } } } return null; } /** * Send a message to the LLM and stream the response */ async sendMessage( content: string, files?: { id: string; name: string; type: string; textContent?: string; preview?: string; }[], enableThinking?: boolean | null, ): Promise { if ((!content.trim() && (!files || files.length === 0)) || this.isLoading) return; if (!this.hasStartedChat) { this.startChat(); } // Capture the target conversation ID at the start of the request const targetConversationId = this.activeConversationId; if (!targetConversationId) return; this.isLoading = true; this.currentResponse = ""; this.ttftMs = null; this.tps = null; this.totalTokens = 0; // Build attachments from files const attachments: MessageAttachment[] = []; let fileContext = ""; if (files && files.length > 0) { for (const file of files) { const isImage = file.type.startsWith("image/"); if (isImage && file.preview) { attachments.push({ type: "image", name: file.name, preview: file.preview, mimeType: file.type, }); } else if (file.textContent) { attachments.push({ type: "text", name: file.name, content: file.textContent, mimeType: file.type, }); // Add text file content to the message context fileContext += `\n\n[File: ${file.name}]\n\`\`\`\n${file.textContent}\n\`\`\``; } else { attachments.push({ type: "file", name: file.name, mimeType: file.type, }); } } } // Combine content with file context const fullContent = content + fileContext; // Add user message directly to the target conversation const userMessage: Message = { id: generateUUID(), role: "user", content: content, // Store original content for display timestamp: Date.now(), attachments: attachments.length > 0 ? attachments : undefined, }; const targetConversation = this.conversations.find( (c) => c.id === targetConversationId, ); if (!targetConversation) { this.isLoading = false; return; } targetConversation.messages.push(userMessage); // Create placeholder for assistant message directly in target conversation const assistantMessage = this.addMessageToConversation( targetConversationId, "assistant", "", ); if (!assistantMessage) { this.isLoading = false; return; } // Sync to this.messages if viewing the target conversation this.syncActiveMessagesIfNeeded(targetConversationId); this.saveConversationsToStorage(); try { // Build the messages array for the API with system prompt const systemPrompt = { role: "system" as const, content: "You are a helpful AI assistant. Respond directly and concisely. Do not show your reasoning or thought process. When files are shared with you, analyze them and respond helpfully.", }; // Build API messages from the target conversation - include file content for text files const apiMessages = [ systemPrompt, ...targetConversation.messages.slice(0, -1).map((m) => { // Build content including any text file attachments let msgContent = m.content; // Add text attachments as context if (m.attachments) { for (const attachment of m.attachments) { if (attachment.type === "text" && attachment.content) { msgContent += `\n\n[File: ${attachment.name}]\n\`\`\`\n${attachment.content}\n\`\`\``; } } } return { role: m.role, content: msgContent, }; }), ]; // Determine the model to use const modelToUse = this.getModelForRequest(); if (!modelToUse) { throw new Error( "No model is loaded yet. Select a model from the sidebar to get started — it will download and load automatically.", ); } const conversationModelInfo = this.buildConversationModelInfo(modelToUse); this.applyConversationModelInfo(conversationModelInfo); // Start timing for TTFT measurement const requestStartTime = performance.now(); let firstTokenTime: number | null = null; let tokenCount = 0; const abortController = new AbortController(); this.currentAbortController = abortController; const response = await fetch("/v1/chat/completions", { method: "POST", headers: { "Content-Type": "application/json", }, body: JSON.stringify({ model: modelToUse, messages: apiMessages, temperature: 0.7, stream: true, logprobs: true, top_logprobs: 5, ...(enableThinking != null && { enable_thinking: enableThinking, }), }), signal: abortController.signal, }); if (!response.ok) { const errorText = await response.text(); throw new Error(`API error: ${response.status} - ${errorText}`); } const reader = response.body?.getReader(); if (!reader) { throw new Error("No response body"); } let streamedContent = ""; let streamedThinking = ""; interface ChatCompletionChunk { choices?: Array<{ delta?: { content?: string; reasoning_content?: string }; logprobs?: { content?: Array<{ token: string; logprob: number; top_logprobs?: Array<{ token: string; logprob: number; bytes: number[] | null; }>; }>; }; }>; } const collectedTokens: TokenData[] = []; await this.parseSSEStream( reader, targetConversationId, (parsed) => { // Clear prefill progress when first token data arrives if (this.prefillProgress) { this.prefillProgress = null; } const choice = parsed.choices?.[0]; const tokenContent = choice?.delta?.content; const thinkingContent = choice?.delta?.reasoning_content; // Collect logprobs data const logprobsContent = choice?.logprobs?.content; if (logprobsContent) { for (const item of logprobsContent) { collectedTokens.push({ token: item.token, logprob: item.logprob, probability: Math.exp(item.logprob), topLogprobs: (item.top_logprobs || []).map((t) => ({ token: t.token, logprob: t.logprob, bytes: t.bytes, })), }); } } if (thinkingContent) { streamedThinking += thinkingContent; } if (tokenContent || thinkingContent) { // Track first token for TTFT if (firstTokenTime === null) { firstTokenTime = performance.now(); this.ttftMs = firstTokenTime - requestStartTime; } // Count tokens (each SSE chunk is typically one token) tokenCount += 1; this.totalTokens = tokenCount; // Update real-time TPS during streaming if (firstTokenTime !== null && tokenCount > 1) { const elapsed = performance.now() - firstTokenTime; this.tps = (tokenCount / elapsed) * 1000; } if (tokenContent) { streamedContent += tokenContent; } // Use stripThinkingTags as fallback for any tags still in content const { displayContent, thinkingContent: tagThinking } = this.stripThinkingTags(streamedContent); const combinedThinking = [streamedThinking, tagThinking] .filter(Boolean) .join("\n\n"); // Only update currentResponse if target conversation is active if (this.activeConversationId === targetConversationId) { this.currentResponse = displayContent; } // Update the assistant message in the target conversation this.updateConversationMessage( targetConversationId, assistantMessage.id, (msg) => { msg.content = displayContent; msg.thinking = combinedThinking || undefined; msg.tokens = [...collectedTokens]; }, ); this.syncActiveMessagesIfNeeded(targetConversationId); this.persistConversation(targetConversationId); } }, { prefill_progress: (data) => { // TaggedModel wraps as {"PrefillProgressChunk": {...}} // model_dump_json() uses snake_case (by_alias defaults to False) const raw = data as Record; const inner = (raw["PrefillProgressChunk"] ?? raw) as { processed_tokens: number; total_tokens: number; }; this.prefillProgress = { processed: inner.processed_tokens, total: inner.total_tokens, startedAt: this.prefillProgress?.startedAt ?? performance.now(), }; }, }, ); // Clear prefill progress after stream ends this.prefillProgress = null; // Calculate final TPS if (firstTokenTime !== null && tokenCount > 1) { const totalGenerationTime = performance.now() - firstTokenTime; this.tps = (tokenCount / totalGenerationTime) * 1000; // tokens per second } // Final cleanup of the message (if conversation still exists) if (this.conversationExists(targetConversationId)) { const { displayContent, thinkingContent: tagThinking } = this.stripThinkingTags(streamedContent); const finalThinking = [streamedThinking, tagThinking] .filter(Boolean) .join("\n\n"); this.updateConversationMessage( targetConversationId, assistantMessage.id, (msg) => { msg.content = displayContent; msg.thinking = finalThinking || undefined; msg.tokens = [...collectedTokens]; // Store performance metrics on the message if (this.ttftMs !== null) { msg.ttftMs = this.ttftMs; } if (this.tps !== null) { msg.tps = this.tps; } }, ); this.syncActiveMessagesIfNeeded(targetConversationId); this.persistConversation(targetConversationId); } } catch (error) { if (error instanceof DOMException && error.name === "AbortError") { // User stopped generation — not an error } else { console.error("Error sending message:", error); this.handleStreamingError( error, targetConversationId, assistantMessage.id, "Failed to get response", ); } } finally { this.currentAbortController = null; this.prefillProgress = null; this.isLoading = false; this.currentResponse = ""; this.saveConversationsToStorage(); } } stopGeneration(): void { this.currentAbortController?.abort(); this.currentAbortController = null; } /** * Generate an image using the image generation API */ async generateImage(prompt: string, modelId?: string): Promise { if (!prompt.trim() || this.isLoading) return; if (!this.hasStartedChat) { this.startChat(); } // Capture the target conversation ID at the start of the request const targetConversationId = this.activeConversationId; if (!targetConversationId) return; this.isLoading = true; this.currentResponse = ""; // Add user message directly to the target conversation const userMessage: Message = { id: generateUUID(), role: "user", content: prompt, timestamp: Date.now(), requestType: "image-generation", }; const targetConversation = this.conversations.find( (c) => c.id === targetConversationId, ); if (!targetConversation) { this.isLoading = false; return; } targetConversation.messages.push(userMessage); // Create placeholder for assistant message directly in target conversation const assistantMessage = this.addMessageToConversation( targetConversationId, "assistant", "Generating image...", ); if (!assistantMessage) { this.isLoading = false; return; } // Sync to this.messages if viewing the target conversation this.syncActiveMessagesIfNeeded(targetConversationId); this.saveConversationsToStorage(); try { // Determine the model to use const model = this.getModelForRequest(modelId); if (!model) { throw new Error( "No model selected. Please select an image generation model.", ); } // Build request body using image generation params const params = this.imageGenerationParams; const hasAdvancedParams = params.seed !== null || params.numInferenceSteps !== null || params.guidance !== null || (params.negativePrompt !== null && params.negativePrompt.trim() !== "") || params.numSyncSteps !== null; const requestBody: Record = { model, prompt, n: params.numImages, quality: params.quality, size: params.size, output_format: params.outputFormat, response_format: "b64_json", stream: params.stream, partial_images: params.partialImages, }; if (hasAdvancedParams) { requestBody.advanced_params = { ...(params.seed !== null && { seed: params.seed }), ...(params.numInferenceSteps !== null && { num_inference_steps: params.numInferenceSteps, }), ...(params.guidance !== null && { guidance: params.guidance }), ...(params.negativePrompt !== null && params.negativePrompt.trim() !== "" && { negative_prompt: params.negativePrompt, }), ...(params.numSyncSteps !== null && { num_sync_steps: params.numSyncSteps, }), }; } const response = await fetch("/v1/images/generations", { method: "POST", headers: { "Content-Type": "application/json", }, body: JSON.stringify(requestBody), }); if (!response.ok) { const errorText = await response.text(); throw new Error(`API error: ${response.status} - ${errorText}`); } // Streaming requires both stream=true AND partialImages > 0 const isStreaming = params.stream && params.partialImages > 0; if (!isStreaming) { // Non-streaming: parse JSON response directly const jsonResponse = (await response.json()) as ImageApiResponse; const format = params.outputFormat || "png"; const mimeType = `image/${format}`; const attachments: MessageAttachment[] = jsonResponse.data .filter((img) => img.b64_json) .map((img, index) => ({ type: "generated-image" as const, name: `generated-image-${index + 1}.${format}`, preview: `data:${mimeType};base64,${img.b64_json}`, mimeType, })); this.updateConversationMessage( targetConversationId, assistantMessage.id, (msg) => { msg.content = ""; msg.attachments = attachments; }, ); this.syncActiveMessagesIfNeeded(targetConversationId); } else { // Streaming mode: use SSE parser const reader = response.body?.getReader(); if (!reader) { throw new Error("No response body"); } interface ImageGenerationChunk { data?: { b64_json?: string }; format?: string; type?: "partial" | "final"; image_index?: number; partial_index?: number; total_partials?: number; } const numImages = params.numImages; await this.parseSSEStream( reader, targetConversationId, (parsed) => { const imageData = parsed.data?.b64_json; if (imageData) { const format = parsed.format || "png"; const mimeType = `image/${format}`; const imageIndex = parsed.image_index ?? 0; if (parsed.type === "partial") { // Update with partial image and progress const partialNum = (parsed.partial_index ?? 0) + 1; const totalPartials = parsed.total_partials ?? 3; const progressText = numImages > 1 ? `Generating image ${imageIndex + 1}/${numImages}... ${partialNum}/${totalPartials}` : `Generating... ${partialNum}/${totalPartials}`; const partialAttachment: MessageAttachment = { type: "generated-image", name: `generated-image.${format}`, preview: `data:${mimeType};base64,${imageData}`, mimeType, }; this.updateConversationMessage( targetConversationId, assistantMessage.id, (msg) => { msg.content = progressText; if (imageIndex === 0) { // First image - safe to replace attachments with partial preview msg.attachments = [partialAttachment]; } else { // Subsequent images - keep existing finals, show partial at current position const existingAttachments = msg.attachments || []; // Keep only the completed final images (up to current imageIndex) const finals = existingAttachments.slice(0, imageIndex); msg.attachments = [...finals, partialAttachment]; } }, ); } else if (parsed.type === "final") { // Final image - replace partial at this position const newAttachment: MessageAttachment = { type: "generated-image", name: `generated-image-${imageIndex + 1}.${format}`, preview: `data:${mimeType};base64,${imageData}`, mimeType, }; this.updateConversationMessage( targetConversationId, assistantMessage.id, (msg) => { if (imageIndex === 0) { // First final image - replace any partial preview msg.attachments = [newAttachment]; } else { // Subsequent images - keep previous finals, replace partial at current position const existingAttachments = msg.attachments || []; // Slice keeps indices 0 to imageIndex-1 (the previous final images) const previousFinals = existingAttachments.slice( 0, imageIndex, ); msg.attachments = [...previousFinals, newAttachment]; } // Update progress message for multiple images if (numImages > 1 && imageIndex < numImages - 1) { msg.content = `Generating image ${imageIndex + 2}/${numImages}...`; } else { msg.content = ""; } }, ); } this.syncActiveMessagesIfNeeded(targetConversationId); } }, ); } } catch (error) { console.error("Error generating image:", error); this.handleStreamingError( error, targetConversationId, assistantMessage.id, "Failed to generate image", ); } finally { this.isLoading = false; this.saveConversationsToStorage(); } } /** * Edit an image using the image edit API */ async editImage( prompt: string, imageDataUrl: string, modelId?: string, ): Promise { if (!prompt.trim() || !imageDataUrl || this.isLoading) return; if (!this.hasStartedChat) { this.startChat(); } // Capture the target conversation ID at the start of the request const targetConversationId = this.activeConversationId; if (!targetConversationId) return; this.isLoading = true; this.currentResponse = ""; // Add user message directly to the target conversation const userMessage: Message = { id: generateUUID(), role: "user", content: prompt, timestamp: Date.now(), requestType: "image-editing", sourceImageDataUrl: imageDataUrl, }; const targetConversation = this.conversations.find( (c) => c.id === targetConversationId, ); if (!targetConversation) { this.isLoading = false; return; } targetConversation.messages.push(userMessage); // Create placeholder for assistant message directly in target conversation const assistantMessage = this.addMessageToConversation( targetConversationId, "assistant", "Editing image...", ); if (!assistantMessage) { this.isLoading = false; return; } // Sync to this.messages if viewing the target conversation this.syncActiveMessagesIfNeeded(targetConversationId); this.saveConversationsToStorage(); // Clear editing state this.editingImage = null; try { // Determine the model to use const model = this.getModelForRequest(modelId); if (!model) { throw new Error( "No model selected. Please select an image generation model.", ); } // Convert base64 data URL to blob const response = await fetch(imageDataUrl); const imageBlob = await response.blob(); // Build FormData request const formData = new FormData(); formData.append("model", model); formData.append("prompt", prompt); formData.append("image", imageBlob, "image.png"); // Add params from image generation params const params = this.imageGenerationParams; formData.append("quality", params.quality); formData.append("size", params.size); formData.append("output_format", params.outputFormat); formData.append("response_format", "b64_json"); formData.append("stream", params.stream ? "1" : "0"); formData.append("partial_images", params.partialImages.toString()); formData.append("input_fidelity", params.inputFidelity); // Advanced params const hasAdvancedParams = params.seed !== null || params.numInferenceSteps !== null || params.guidance !== null || (params.negativePrompt !== null && params.negativePrompt.trim() !== "") || params.numSyncSteps !== null; if (hasAdvancedParams) { formData.append( "advanced_params", JSON.stringify({ ...(params.seed !== null && { seed: params.seed }), ...(params.numInferenceSteps !== null && { num_inference_steps: params.numInferenceSteps, }), ...(params.guidance !== null && { guidance: params.guidance }), ...(params.negativePrompt !== null && params.negativePrompt.trim() !== "" && { negative_prompt: params.negativePrompt, }), ...(params.numSyncSteps !== null && { num_sync_steps: params.numSyncSteps, }), }), ); } const apiResponse = await fetch("/v1/images/edits", { method: "POST", body: formData, }); if (!apiResponse.ok) { const errorText = await apiResponse.text(); throw new Error(`API error: ${apiResponse.status} - ${errorText}`); } // Streaming requires both stream=true AND partialImages > 0 const isStreaming = params.stream && params.partialImages > 0; if (!isStreaming) { // Non-streaming: parse JSON response directly const jsonResponse = (await apiResponse.json()) as ImageApiResponse; const format = params.outputFormat || "png"; const mimeType = `image/${format}`; const attachments: MessageAttachment[] = jsonResponse.data .filter((img) => img.b64_json) .map((img) => ({ type: "generated-image" as const, name: `edited-image.${format}`, preview: `data:${mimeType};base64,${img.b64_json}`, mimeType, })); this.updateConversationMessage( targetConversationId, assistantMessage.id, (msg) => { msg.content = ""; msg.attachments = attachments; }, ); this.syncActiveMessagesIfNeeded(targetConversationId); } else { // Streaming mode: use SSE parser const reader = apiResponse.body?.getReader(); if (!reader) { throw new Error("No response body"); } interface ImageEditChunk { data?: { b64_json?: string }; format?: string; type?: "partial" | "final"; partial_index?: number; total_partials?: number; } await this.parseSSEStream( reader, targetConversationId, (parsed) => { const imageData = parsed.data?.b64_json; if (imageData) { const format = parsed.format || "png"; const mimeType = `image/${format}`; if (parsed.type === "partial") { // Update with partial image and progress const partialNum = (parsed.partial_index ?? 0) + 1; const totalPartials = parsed.total_partials ?? 3; this.updateConversationMessage( targetConversationId, assistantMessage.id, (msg) => { msg.content = `Editing... ${partialNum}/${totalPartials}`; msg.attachments = [ { type: "generated-image", name: `edited-image.${format}`, preview: `data:${mimeType};base64,${imageData}`, mimeType, }, ]; }, ); } else if (parsed.type === "final") { // Final image this.updateConversationMessage( targetConversationId, assistantMessage.id, (msg) => { msg.content = ""; msg.attachments = [ { type: "generated-image", name: `edited-image.${format}`, preview: `data:${mimeType};base64,${imageData}`, mimeType, }, ]; }, ); } this.syncActiveMessagesIfNeeded(targetConversationId); } }, ); } } catch (error) { console.error("Error editing image:", error); this.handleStreamingError( error, targetConversationId, assistantMessage.id, "Failed to edit image", ); } finally { this.isLoading = false; this.saveConversationsToStorage(); } } /** * Clear current chat and go back to welcome state */ clearChat() { this.activeConversationId = null; this.messages = []; this.hasStartedChat = false; this.isTopologyMinimized = false; this.currentResponse = ""; // Clear performance stats this.ttftMs = null; this.tps = null; } /** * Get the active conversation */ getActiveConversation(): Conversation | null { if (!this.activeConversationId) return null; return ( this.conversations.find((c) => c.id === this.activeConversationId) || null ); } /** * Update the thinking preference for the active conversation */ setConversationThinking(enabled: boolean) { this.thinkingEnabled = enabled; const conv = this.getActiveConversation(); if (conv) { conv.enableThinking = enabled; this.saveConversationsToStorage(); } } /** * Start a download on a specific node */ async startDownload(nodeId: string, shardMetadata: object): Promise { try { const response = await fetch("/download/start", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ targetNodeId: nodeId, shardMetadata: shardMetadata, }), }); if (!response.ok) { const errorText = await response.text(); throw new Error( `Failed to start download: ${response.status} - ${errorText}`, ); } } catch (error) { console.error("Error starting download:", error); throw error; } } /** * Delete a downloaded model from a specific node */ async deleteDownload(nodeId: string, modelId: string): Promise { try { const response = await fetch( `/download/${encodeURIComponent(nodeId)}/${encodeURIComponent(modelId)}`, { method: "DELETE", }, ); if (!response.ok) { const errorText = await response.text(); throw new Error( `Failed to delete download: ${response.status} - ${errorText}`, ); } } catch (error) { console.error("Error deleting download:", error); throw error; } } /** * List all available traces */ async listTraces(): Promise { const response = await fetch("/v1/traces"); if (!response.ok) { throw new Error(`Failed to list traces: ${response.status}`); } return (await response.json()) as TraceListResponse; } /** * Check if a trace exists for a given task ID */ async checkTraceExists(taskId: string): Promise { try { const response = await fetch(`/v1/traces/${encodeURIComponent(taskId)}`); return response.ok; } catch { return false; } } /** * Get computed statistics for a task's trace */ async fetchTraceStats(taskId: string): Promise { const response = await fetch( `/v1/traces/${encodeURIComponent(taskId)}/stats`, ); if (!response.ok) { throw new Error(`Failed to fetch trace stats: ${response.status}`); } return (await response.json()) as TraceStatsResponse; } /** * Delete traces by task IDs */ async deleteTraces( taskIds: string[], ): Promise<{ deleted: string[]; notFound: string[] }> { const response = await fetch("/v1/traces/delete", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ taskIds }), }); if (!response.ok) { throw new Error(`Failed to delete traces: ${response.status}`); } return await response.json(); } /** * Get the URL for the raw trace file (for Perfetto) */ getTraceRawUrl(taskId: string): string { return `/v1/traces/${encodeURIComponent(taskId)}/raw`; } } export const appStore = new AppStore(); // Reactive exports export const hasStartedChat = () => appStore.hasStartedChat; export const messages = () => appStore.messages; export const currentResponse = () => appStore.currentResponse; export const isLoading = () => appStore.isLoading; export const ttftMs = () => appStore.ttftMs; export const tps = () => appStore.tps; export const totalTokens = () => appStore.totalTokens; export const prefillProgress = () => appStore.prefillProgress; export const topologyData = () => appStore.topologyData; export const instances = () => appStore.instances; export const runners = () => appStore.runners; export const downloads = () => appStore.downloads; export const nodeDisk = () => appStore.nodeDisk; export const placementPreviews = () => appStore.placementPreviews; export const selectedPreviewModelId = () => appStore.selectedPreviewModelId; export const isLoadingPreviews = () => appStore.isLoadingPreviews; export const lastUpdate = () => appStore.lastUpdate; export const isTopologyMinimized = () => appStore.isTopologyMinimized; export const selectedChatModel = () => appStore.selectedChatModel; export const thinkingEnabled = () => appStore.thinkingEnabled; export const debugMode = () => appStore.getDebugMode(); export const topologyOnlyMode = () => appStore.getTopologyOnlyMode(); export const chatSidebarVisible = () => appStore.getChatSidebarVisible(); // Actions export const stopGeneration = () => appStore.stopGeneration(); export const startChat = () => appStore.startChat(); export const sendMessage = ( content: string, files?: { id: string; name: string; type: string; textContent?: string; preview?: string; }[], enableThinking?: boolean | null, ) => appStore.sendMessage(content, files, enableThinking); export const generateImage = (prompt: string, modelId?: string) => appStore.generateImage(prompt, modelId); export const editImage = ( prompt: string, imageDataUrl: string, modelId?: string, ) => appStore.editImage(prompt, imageDataUrl, modelId); export const editingImage = () => appStore.editingImage; export const setEditingImage = (imageDataUrl: string, sourceMessage: Message) => appStore.setEditingImage(imageDataUrl, sourceMessage); export const clearEditingImage = () => appStore.clearEditingImage(); export const clearChat = () => appStore.clearChat(); export const setSelectedChatModel = (modelId: string) => appStore.setSelectedModel(modelId); export const selectPreviewModel = (modelId: string | null) => appStore.selectPreviewModel(modelId); export const togglePreviewNodeFilter = (nodeId: string) => appStore.togglePreviewNodeFilter(nodeId); export const clearPreviewNodeFilter = () => appStore.clearPreviewNodeFilter(); export const previewNodeFilter = () => appStore.previewNodeFilter; export const deleteMessage = (messageId: string) => appStore.deleteMessage(messageId); export const editMessage = (messageId: string, newContent: string) => appStore.editMessage(messageId, newContent); export const editAndRegenerate = (messageId: string, newContent: string) => appStore.editAndRegenerate(messageId, newContent); export const regenerateLastResponse = () => appStore.regenerateLastResponse(); export const regenerateFromToken = (messageId: string, tokenIndex: number) => appStore.regenerateFromToken(messageId, tokenIndex); // Conversation actions export const conversations = () => appStore.conversations; export const activeConversationId = () => appStore.activeConversationId; export const createConversation = (name?: string) => appStore.createConversation(name); export const loadConversation = (id: string) => appStore.loadConversation(id); export const deleteConversation = (id: string) => appStore.deleteConversation(id); export const deleteAllConversations = () => appStore.deleteAllConversations(); export const renameConversation = (id: string, name: string) => appStore.renameConversation(id, name); export const getActiveConversation = () => appStore.getActiveConversation(); export const setConversationThinking = (enabled: boolean) => appStore.setConversationThinking(enabled); // Sidebar actions export const isSidebarOpen = () => appStore.isSidebarOpen; export const toggleSidebar = () => appStore.toggleSidebar(); export const toggleDebugMode = () => appStore.toggleDebugMode(); export const setDebugMode = (enabled: boolean) => appStore.setDebugMode(enabled); export const toggleTopologyOnlyMode = () => appStore.toggleTopologyOnlyMode(); export const setTopologyOnlyMode = (enabled: boolean) => appStore.setTopologyOnlyMode(enabled); export const toggleChatSidebarVisible = () => appStore.toggleChatSidebarVisible(); export const setChatSidebarVisible = (visible: boolean) => appStore.setChatSidebarVisible(visible); // Mobile sidebar state export const mobileChatSidebarOpen = () => appStore.mobileChatSidebarOpen; export const toggleMobileChatSidebar = () => appStore.toggleMobileChatSidebar(); export const setMobileChatSidebarOpen = (open: boolean) => appStore.setMobileChatSidebarOpen(open); export const mobileRightSidebarOpen = () => appStore.mobileRightSidebarOpen; export const toggleMobileRightSidebar = () => appStore.toggleMobileRightSidebar(); export const setMobileRightSidebarOpen = (open: boolean) => appStore.setMobileRightSidebarOpen(open); export const refreshState = () => appStore.fetchState(); // Connection status export const isConnected = () => appStore.isConnected; // Node identities (for OS version mismatch detection) export const nodeIdentities = () => appStore.nodeIdentities; // Thunderbolt & RDMA status export const nodeThunderbolt = () => appStore.nodeThunderbolt; export const nodeRdmaCtl = () => appStore.nodeRdmaCtl; export const thunderboltBridgeCycles = () => appStore.thunderboltBridgeCycles; export const nodeThunderboltBridge = () => appStore.nodeThunderboltBridge; // Image generation params export const imageGenerationParams = () => appStore.getImageGenerationParams(); export const setImageGenerationParams = ( params: Partial, ) => appStore.setImageGenerationParams(params); export const resetImageGenerationParams = () => appStore.resetImageGenerationParams(); // Download actions export const startDownload = (nodeId: string, shardMetadata: object) => appStore.startDownload(nodeId, shardMetadata); export const deleteDownload = (nodeId: string, modelId: string) => appStore.deleteDownload(nodeId, modelId); // Trace actions export const listTraces = () => appStore.listTraces(); export const checkTraceExists = (taskId: string) => appStore.checkTraceExists(taskId); export const fetchTraceStats = (taskId: string) => appStore.fetchTraceStats(taskId); export const getTraceRawUrl = (taskId: string) => appStore.getTraceRawUrl(taskId); export const deleteTraces = (taskIds: string[]) => appStore.deleteTraces(taskIds); ================================================ FILE: dashboard/src/lib/stores/favorites.svelte.ts ================================================ /** * FavoritesStore - Manages favorite models with localStorage persistence */ import { browser } from "$app/environment"; const FAVORITES_KEY = "exo-favorite-models"; class FavoritesStore { favorites = $state>(new Set()); constructor() { if (browser) { this.loadFromStorage(); } } private loadFromStorage() { try { const stored = localStorage.getItem(FAVORITES_KEY); if (stored) { const parsed = JSON.parse(stored) as string[]; this.favorites = new Set(parsed); } } catch (error) { console.error("Failed to load favorites:", error); } } private saveToStorage() { try { const array = Array.from(this.favorites); localStorage.setItem(FAVORITES_KEY, JSON.stringify(array)); } catch (error) { console.error("Failed to save favorites:", error); } } add(baseModelId: string) { const next = new Set(this.favorites); next.add(baseModelId); this.favorites = next; this.saveToStorage(); } remove(baseModelId: string) { const next = new Set(this.favorites); next.delete(baseModelId); this.favorites = next; this.saveToStorage(); } toggle(baseModelId: string) { if (this.favorites.has(baseModelId)) { this.remove(baseModelId); } else { this.add(baseModelId); } } isFavorite(baseModelId: string): boolean { return this.favorites.has(baseModelId); } getAll(): string[] { return Array.from(this.favorites); } getSet(): Set { return new Set(this.favorites); } hasAny(): boolean { return this.favorites.size > 0; } clearAll() { this.favorites = new Set(); this.saveToStorage(); } } export const favoritesStore = new FavoritesStore(); export const favorites = () => favoritesStore.favorites; export const hasFavorites = () => favoritesStore.hasAny(); export const isFavorite = (baseModelId: string) => favoritesStore.isFavorite(baseModelId); export const toggleFavorite = (baseModelId: string) => favoritesStore.toggle(baseModelId); export const addFavorite = (baseModelId: string) => favoritesStore.add(baseModelId); export const removeFavorite = (baseModelId: string) => favoritesStore.remove(baseModelId); export const getFavorites = () => favoritesStore.getAll(); export const getFavoritesSet = () => favoritesStore.getSet(); export const clearFavorites = () => favoritesStore.clearAll(); ================================================ FILE: dashboard/src/lib/stores/recents.svelte.ts ================================================ /** * RecentsStore - Manages recently launched models with localStorage persistence */ import { browser } from "$app/environment"; const RECENTS_KEY = "exo-recent-models"; const MAX_RECENT_MODELS = 20; interface RecentEntry { modelId: string; launchedAt: number; } class RecentsStore { recents = $state([]); constructor() { if (browser) { this.loadFromStorage(); } } private loadFromStorage() { try { const stored = localStorage.getItem(RECENTS_KEY); if (stored) { const parsed = JSON.parse(stored) as RecentEntry[]; this.recents = parsed; } } catch (error) { console.error("Failed to load recent models:", error); } } private saveToStorage() { try { localStorage.setItem(RECENTS_KEY, JSON.stringify(this.recents)); } catch (error) { console.error("Failed to save recent models:", error); } } recordLaunch(modelId: string) { // Remove existing entry for this model (if any) to move it to top const filtered = this.recents.filter((r) => r.modelId !== modelId); // Prepend new entry const next = [{ modelId, launchedAt: Date.now() }, ...filtered]; // Cap at max this.recents = next.slice(0, MAX_RECENT_MODELS); this.saveToStorage(); } getRecentModelIds(): string[] { return this.recents.map((r) => r.modelId); } hasAny(): boolean { return this.recents.length > 0; } clearAll() { this.recents = []; this.saveToStorage(); } } export const recentsStore = new RecentsStore(); export const hasRecents = () => recentsStore.hasAny(); export const getRecentModelIds = () => recentsStore.getRecentModelIds(); export const getRecentEntries = () => recentsStore.recents; export const recordRecentLaunch = (modelId: string) => recentsStore.recordLaunch(modelId); export const clearRecents = () => recentsStore.clearAll(); ================================================ FILE: dashboard/src/lib/stores/toast.svelte.ts ================================================ /** * Toast notification store - Global notification system for the EXO dashboard. * * Usage: * import { addToast, dismissToast, toasts } from "$lib/stores/toast.svelte"; * addToast({ type: "success", message: "Model launched" }); * addToast({ type: "error", message: "Connection lost", persistent: true }); */ type ToastType = "success" | "error" | "warning" | "info"; export interface Toast { id: string; type: ToastType; message: string; /** Auto-dismiss after this many ms. 0 = persistent (must be dismissed manually). */ duration: number; createdAt: number; } interface ToastInput { type: ToastType; message: string; /** If true, toast stays until manually dismissed. Default: false. */ persistent?: boolean; /** Auto-dismiss duration in ms. Default: 4000 for success/info, 6000 for error/warning. */ duration?: number; } const DEFAULT_DURATIONS: Record = { success: 4000, info: 4000, warning: 6000, error: 6000, }; let toastList = $state([]); const timers = new Map>(); function generateId(): string { return `toast-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; } export function addToast(input: ToastInput): string { const id = generateId(); const duration = input.persistent ? 0 : (input.duration ?? DEFAULT_DURATIONS[input.type]); const toast: Toast = { id, type: input.type, message: input.message, duration, createdAt: Date.now(), }; toastList = [...toastList, toast]; if (duration > 0) { const timer = setTimeout(() => dismissToast(id), duration); timers.set(id, timer); } return id; } export function dismissToast(id: string): void { const timer = timers.get(id); if (timer) { clearTimeout(timer); timers.delete(id); } toastList = toastList.filter((t) => t.id !== id); } /** Dismiss all toasts matching a message (useful for dedup). */ export function dismissByMessage(message: string): void { const matching = toastList.filter((t) => t.message === message); for (const t of matching) { dismissToast(t.id); } } export function toasts(): Toast[] { return toastList; } ================================================ FILE: dashboard/src/lib/types/files.ts ================================================ /** * File attachment types for the chat interface */ export interface ChatUploadedFile { id: string; name: string; size: number; type: string; file: File; preview?: string; textContent?: string; } export interface ChatAttachment { type: "image" | "text" | "pdf" | "audio"; name: string; content?: string; base64Url?: string; mimeType?: string; } export type FileCategory = "image" | "text" | "pdf" | "audio" | "unknown"; export const IMAGE_EXTENSIONS = [ ".jpg", ".jpeg", ".png", ".gif", ".webp", ".svg", ]; export const IMAGE_MIME_TYPES = [ "image/jpeg", "image/png", "image/gif", "image/webp", "image/svg+xml", ]; export const TEXT_EXTENSIONS = [ ".txt", ".md", ".json", ".xml", ".yaml", ".yml", ".csv", ".log", ".js", ".ts", ".jsx", ".tsx", ".py", ".java", ".cpp", ".c", ".h", ".css", ".html", ".htm", ".sql", ".sh", ".bat", ".rs", ".go", ".rb", ".php", ".swift", ".kt", ".scala", ".r", ".dart", ".vue", ".svelte", ]; export const TEXT_MIME_TYPES = [ "text/plain", "text/markdown", "text/csv", "text/html", "text/css", "application/json", "application/xml", "text/xml", "application/javascript", "text/javascript", "application/typescript", ]; export const PDF_EXTENSIONS = [".pdf"]; export const PDF_MIME_TYPES = ["application/pdf"]; export const AUDIO_EXTENSIONS = [".mp3", ".wav", ".ogg", ".m4a"]; export const AUDIO_MIME_TYPES = [ "audio/mpeg", "audio/wav", "audio/ogg", "audio/mp4", ]; /** * Get file category based on MIME type and extension */ export function getFileCategory( mimeType: string, fileName: string, ): FileCategory { const extension = fileName.toLowerCase().slice(fileName.lastIndexOf(".")); if ( IMAGE_MIME_TYPES.includes(mimeType) || IMAGE_EXTENSIONS.includes(extension) ) { return "image"; } if (PDF_MIME_TYPES.includes(mimeType) || PDF_EXTENSIONS.includes(extension)) { return "pdf"; } if ( AUDIO_MIME_TYPES.includes(mimeType) || AUDIO_EXTENSIONS.includes(extension) ) { return "audio"; } if ( TEXT_MIME_TYPES.includes(mimeType) || TEXT_EXTENSIONS.includes(extension) || mimeType.startsWith("text/") ) { return "text"; } return "unknown"; } /** * Get accept string for file input based on categories */ export function getAcceptString(categories: FileCategory[]): string { const accepts: string[] = []; for (const category of categories) { switch (category) { case "image": accepts.push(...IMAGE_EXTENSIONS, ...IMAGE_MIME_TYPES); break; case "text": accepts.push(...TEXT_EXTENSIONS, ...TEXT_MIME_TYPES); break; case "pdf": accepts.push(...PDF_EXTENSIONS, ...PDF_MIME_TYPES); break; case "audio": accepts.push(...AUDIO_EXTENSIONS, ...AUDIO_MIME_TYPES); break; } } return accepts.join(","); } /** * Format file size for display */ export function formatFileSize(bytes: number): string { if (bytes === 0) return "0 B"; const k = 1024; const sizes = ["B", "KB", "MB", "GB"]; const i = Math.floor(Math.log(bytes) / Math.log(k)); return parseFloat((bytes / Math.pow(k, i)).toFixed(1)) + " " + sizes[i]; } /** * Read file as data URL (base64) */ export function readFileAsDataURL(file: File): Promise { return new Promise((resolve, reject) => { const reader = new FileReader(); reader.onload = () => resolve(reader.result as string); reader.onerror = () => reject(reader.error); reader.readAsDataURL(file); }); } /** * Read file as text */ export function readFileAsText(file: File): Promise { return new Promise((resolve, reject) => { const reader = new FileReader(); reader.onload = () => resolve(reader.result as string); reader.onerror = () => reject(reader.error); reader.readAsText(file); }); } /** * Process uploaded files into ChatUploadedFile format */ export async function processUploadedFiles( files: File[], ): Promise { const results: ChatUploadedFile[] = []; for (const file of files) { const id = Date.now().toString() + Math.random().toString(36).substring(2, 9); const category = getFileCategory(file.type, file.name); const base: ChatUploadedFile = { id, name: file.name, size: file.size, type: file.type, file, }; try { if (category === "image") { const preview = await readFileAsDataURL(file); results.push({ ...base, preview }); } else if (category === "text" || category === "unknown") { const textContent = await readFileAsText(file); results.push({ ...base, textContent }); } else if (category === "pdf") { results.push(base); } else if (category === "audio") { const preview = await readFileAsDataURL(file); results.push({ ...base, preview }); } else { results.push(base); } } catch (error) { console.error("Error processing file:", file.name, error); results.push(base); } } return results; } ================================================ FILE: dashboard/src/lib/utils/downloads.ts ================================================ /** * Shared utilities for parsing and querying download state. * * The download state from `/state` is shaped as: * Record> * * Each entry is a tagged union object like: * { "DownloadCompleted": { shard_metadata: { "PipelineShardMetadata": { model_card: { model_id: "..." }, ... } }, ... } } */ /** Unwrap one level of tagged-union envelope, returning [tag, payload]. */ function unwrapTagged( obj: Record, ): [string, Record] | null { const keys = Object.keys(obj); if (keys.length !== 1) return null; const tag = keys[0]; const payload = obj[tag]; if (!payload || typeof payload !== "object") return null; return [tag, payload as Record]; } /** Extract the model ID string from a download entry's nested shard_metadata. */ export function extractModelIdFromDownload( downloadPayload: Record, ): string | null { const shardMetadata = downloadPayload.shard_metadata ?? downloadPayload.shardMetadata; if (!shardMetadata || typeof shardMetadata !== "object") return null; const unwrapped = unwrapTagged(shardMetadata as Record); if (!unwrapped) return null; const [, shardData] = unwrapped; const modelMeta = shardData.model_card ?? shardData.modelCard; if (!modelMeta || typeof modelMeta !== "object") return null; const meta = modelMeta as Record; return (meta.model_id as string) ?? (meta.modelId as string) ?? null; } /** Extract the shard_metadata object from a download entry payload. */ export function extractShardMetadata( downloadPayload: Record, ): Record | null { const shardMetadata = downloadPayload.shard_metadata ?? downloadPayload.shardMetadata; if (!shardMetadata || typeof shardMetadata !== "object") return null; return shardMetadata as Record; } /** Get the download tag (DownloadCompleted, DownloadOngoing, etc.) from a wrapped entry. */ export function getDownloadTag( entry: unknown, ): [string, Record] | null { if (!entry || typeof entry !== "object") return null; return unwrapTagged(entry as Record); } /** * Iterate over all download entries for a given node, yielding [tag, payload, modelId]. */ function* iterNodeDownloads( nodeDownloads: unknown[], ): Generator<[string, Record, string]> { for (const entry of nodeDownloads) { const tagged = getDownloadTag(entry); if (!tagged) continue; const [tag, payload] = tagged; const modelId = extractModelIdFromDownload(payload); if (!modelId) continue; yield [tag, payload, modelId]; } } /** Check if a specific model is fully downloaded (DownloadCompleted) on a specific node. */ export function isModelDownloadedOnNode( downloadsData: Record, nodeId: string, modelId: string, ): boolean { const nodeDownloads = downloadsData[nodeId]; if (!Array.isArray(nodeDownloads)) return false; for (const [tag, , entryModelId] of iterNodeDownloads(nodeDownloads)) { if (tag === "DownloadCompleted" && entryModelId === modelId) return true; } return false; } /** Get all node IDs where a model is fully downloaded (DownloadCompleted). */ export function getNodesWithModelDownloaded( downloadsData: Record, modelId: string, ): string[] { const result: string[] = []; for (const nodeId of Object.keys(downloadsData)) { if (isModelDownloadedOnNode(downloadsData, nodeId, modelId)) { result.push(nodeId); } } return result; } /** * Find shard metadata for a model from any download entry across all nodes. * Returns the first match found (completed entries are preferred). */ export function getShardMetadataForModel( downloadsData: Record, modelId: string, ): Record | null { let fallback: Record | null = null; for (const nodeDownloads of Object.values(downloadsData)) { if (!Array.isArray(nodeDownloads)) continue; for (const [tag, payload, entryModelId] of iterNodeDownloads( nodeDownloads, )) { if (entryModelId !== modelId) continue; const shard = extractShardMetadata(payload); if (!shard) continue; if (tag === "DownloadCompleted") return shard; if (!fallback) fallback = shard; } } return fallback; } /** * Get the download status tag for a specific model on a specific node. * Returns the "best" status: DownloadCompleted > DownloadOngoing > others. */ export function getModelDownloadStatus( downloadsData: Record, nodeId: string, modelId: string, ): string | null { const nodeDownloads = downloadsData[nodeId]; if (!Array.isArray(nodeDownloads)) return null; let best: string | null = null; for (const [tag, , entryModelId] of iterNodeDownloads(nodeDownloads)) { if (entryModelId !== modelId) continue; if (tag === "DownloadCompleted") return tag; if (tag === "DownloadOngoing") best = tag; else if (!best) best = tag; } return best; } ================================================ FILE: dashboard/src/routes/+layout.svelte ================================================ EXO
{@render children?.()}
================================================ FILE: dashboard/src/routes/+page.svelte ================================================ {#snippet clusterWarnings()} {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed) || (macStudioEn2RdmaWarning && !macStudioEn2Dismissed)}
{#if tbBridgeCycles.length > 0} {@const cycle = tbBridgeCycles[0]} {@const serviceName = getTbBridgeServiceName(cycle)} {@const disableCmd = `sudo networksetup -setnetworkserviceenabled "${serviceName}" off`} {/if} {#if macosVersionMismatch} {/if} {#if tb5WithoutRdma && !tb5InfoDismissed}
RDMA NOT ENABLED
{/if} {#if macStudioEn2RdmaWarning && !macStudioEn2Dismissed} {/if}
{/if} {/snippet} {#snippet clusterWarningsCompact()} {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed) || (macStudioEn2RdmaWarning && !macStudioEn2Dismissed)}
{#if tbBridgeCycles.length > 0}
TB CYCLE
{/if} {#if macosVersionMismatch}
macOS MISMATCH
{/if} {#if tb5WithoutRdma && !tb5InfoDismissed}
RDMA NOT ENABLED
{/if} {#if macStudioEn2RdmaWarning && !macStudioEn2Dismissed}
BAD RDMA PORT
{/if}
{/if} {/snippet} { if (!showOnboardingOverlay || stepTransitioning) return; if (e.key === "ArrowRight" || e.key === " " || e.key === "Enter") { if (onboardingStep >= 1 && onboardingStep <= 4 && showContinueButton) { e.preventDefault(); advanceStep(onboardingStep < 4 ? onboardingStep + 1 : 6); } } }} />
{#if showOnboardingOverlay}
{#if onboardingStep >= 1 && onboardingStep <= 4}
exo

{onboardingStep === 1 ? "EXO connects all your devices into an AI supercomputer." : stepTitle}

{#if onboardingStep === 2}   {:else if onboardingStep === 3} The model is automatically distributed. Each device handles a piece. {:else if onboardingStep === 4} {stepTitle === "exo self-heals" ? "exo automatically redistributes the model so inference continues without interruption." : "Devices can leave anytime. Laptops close, machines restart."} {:else}   {/if}

Your EXO Network

{#if onboardingStep <= 1 || $topologyOpacity > 0.01}
{/if} {userDeviceInfo.name} {userDeviceInfo.memoryGB}{" "}GB Mac Studio {SIMULATED_STUDIO_GB}{" "}GB (example) 0.5 ? "rgba(220,38,38,0.7)" : "rgba(255,255,255,0.15)"} stroke-width="1.5" stroke-dasharray="6,6" opacity={$connectionOpacity} class={$connectionIsRed > 0.5 ? "onboarding-connection-line-red" : "onboarding-connection-line"} /> {#if $disconnectXOpacity > 0.01} {/if} {onboardingCombinedGB} GB combined {#if unlockedModels.length > 0 && $chipPhase > 0.01} {@const centerX = ($device1X + $device2X) / 2} {@const chipW = 140} {@const chipH = 30} {@const chipGap = 12} {@const totalW = unlockedModels.length * chipW + (unlockedModels.length - 1) * chipGap} {@const startX = centerX - totalW / 2} {@const headerProgress = Math.min(1, $chipPhase)} {@const headerY = 332 + 12 * (1 - headerProgress)} {@const yellowR = 234} {@const yellowG = 179} {@const yellowB = 8} NEW MODELS UNLOCKED {#each unlockedModels as model, i} {@const stagger = i * 0.6} {@const progress = Math.max( 0, Math.min(1, $chipPhase - stagger), )} {@const modelName = ( model.name || model.id.split("/").pop() || "" ).slice(0, 18)} {@const modelSize = Math.round(getModelSizeGB(model))} {@const slideY = 16 * (1 - progress)} {@const chipScale = 0.85 + 0.15 * progress} {@const highlightPeak = progress < 0.6 ? progress / 0.6 : 1 - ((progress - 0.6) / 0.4) * 0.6} {@const borderYellow = 0.15 + 0.35 * highlightPeak} {@const fillYellow = 0.02 + 0.06 * highlightPeak} {#if progress > 0} 0.3 ? "url(#chip-glow)" : "none"} > {modelName} {#if modelSize} {modelSize} GB {/if} {/if} {/each} {/if} {#if $modelBlockOpacity > 0.01} {#if $modelSplitProgress < 0.05} {@const modelCenterX = $device2Opacity > 0.3 ? ($device1X + $device2X) / 2 : $device1X} LLM {:else} {@const splitX = $modelSplitProgress * (($device2X - $device1X) / 2)} {@const centerX = ($device1X + $device2X) / 2} {@const splitY = $modelBlockY + $modelSplitProgress * 80} Shard 1/2 Shard 2/2 {/if} {/if}
{:else if onboardingStep === 6}

Choose a model

Showing recommended models for your devices ({Math.round( clusterMemory().total / (1024 * 1024 * 1024), )} GB memory available).

{#if onboardingError}
{onboardingError}
{/if} {#if onboardingModels.length === 0}
Loading models...
{:else}
{#each onboardingModels as model} {@const sizeGB = getModelSizeGB(model)} {@const fitsNow = hasEnoughMemory(model)} {@const tags = modelTags()[model.id] || []} {/each}
{/if}
{:else if onboardingStep === 7}

Downloading

{#if onboardingModelId}

{onboardingModelId.split("/").pop() ?? onboardingModelId}

{/if}
{#if onboardingDownloadProgress}
{onboardingDownloadProgress.percentage.toFixed(1)}% {formatBytes(onboardingDownloadProgress.downloadedBytes)} / {formatBytes(onboardingDownloadProgress.totalBytes)}
{formatSpeed(onboardingDownloadProgress.speed)} ETA: {formatEta(onboardingDownloadProgress.etaMs)}
{:else}

Preparing download...

{/if}

This may take a few minutes depending on your connection.

{:else if onboardingStep === 8}

Loading into memory

{#if onboardingModelId}

{onboardingModelId.split("/").pop() ?? onboardingModelId}

{/if}
{#if onboardingLoadProgress}

{onboardingLoadProgress.layersLoaded} / {onboardingLoadProgress.totalLayers} layers loaded

{:else}

Loading...

{/if}
{:else if onboardingStep === 9}
exo {#if onboardingModelId}

{onboardingModelId.split("/").pop() ?? onboardingModelId}

{/if}
{#each suggestedPrompts as chip} {/each}
{/if}
{#if onboardingStep === 6} m.id))} canModelFit={(modelId) => { const model = models.find((m) => m.id === modelId); return model ? hasEnoughMemory(model) : false; }} getModelFitStatus={(modelId): ModelMemoryFitStatus => { const model = models.find((m) => m.id === modelId); return model ? getModelMemoryFitStatus(model) : "too_large"; }} onSelect={(modelId) => { isModelPickerOpen = false; onboardingLaunchModel(modelId); }} onClose={() => (isModelPickerOpen = false)} onToggleFavorite={toggleFavorite} onAddModel={addModelFromPicker} onDeleteModel={deleteCustomModel} totalMemoryGB={clusterMemory().total / (1024 * 1024 * 1024)} usedMemoryGB={clusterMemory().used / (1024 * 1024 * 1024)} {downloadsData} topologyNodes={data?.nodes} /> {/if} {/if} {#if !topologyOnlyEnabled} {/if} {#if !topologyOnlyEnabled} setMobileChatSidebarOpen(false)} onNewChat={handleNewChat} onSelectConversation={() => { userForcedIdle = false; }} /> {/if}
{#if !topologyOnlyEnabled && sidebarVisible} {/if} {#if topologyOnlyEnabled}
{@render clusterWarnings()} {#if tb5WithoutRdma && !tb5InfoDismissed}
0} class:top-4={tbBridgeCycles.length === 0} role="status" >
RDMA NOT ENABLED
{/if}
{:else if !chatStarted}
{#if !update}

Connecting to cluster…

{/if} {@render clusterWarnings()} {#if tb5WithoutRdma && !tb5InfoDismissed}
0} class:top-4={tbBridgeCycles.length === 0} role="status" >
RDMA NOT ENABLED
{/if} {#if isFilterActive()} {/if}
{#if instanceCount === 0}

Select a model to get started.

{/if}
{#if mobileRightOpen} {/if} {#snippet rightSidebarContent()} {#if instanceCount > 0}

Instances

{#each Object.entries(instanceData) as [id, instance]} {@const downloadInfo = getInstanceDownloadStatus( id, instance, )} {@const statusText = downloadInfo.statusText} {@const isDownloading = downloadInfo.isDownloading} {@const isFailed = statusText === "FAILED"} {@const isLoading = statusText === "LOADING"} {@const isWarmingUp = statusText === "WARMING UP" || statusText === "WAITING"} {@const isReady = statusText === "READY" || statusText === "LOADED"} {@const isRunning = statusText === "RUNNING"} {@const instanceModelId = getInstanceModelId(instance)} {@const instanceInfo = getInstanceInfo(instance)} {@const instanceConnections = getInstanceConnections(instance)}
(hoveredInstanceId = id)} onmouseleave={() => (hoveredInstanceId = null)} onclick={() => { if ( instanceModelId && instanceModelId !== "Unknown" && instanceModelId !== "Unknown Model" ) { userForcedIdle = false; setSelectedChatModel(instanceModelId); } }} onkeydown={(e) => { if (e.key === "Enter" || e.key === " ") { if ( instanceModelId && instanceModelId !== "Unknown" && instanceModelId !== "Unknown Model" ) { setSelectedChatModel(instanceModelId); } } }} >
{id.slice(0, 8).toUpperCase()}
{getInstanceModelId(instance)}
{instanceInfo.sharding} · {instanceInfo.instanceType} {statusText}
{#if instanceModelId && instanceModelId !== "Unknown" && instanceModelId !== "Unknown Model"} Hugging Face {/if} {#if instanceInfo.nodeNames.length > 0}
{instanceInfo.nodeNames.join(", ")}
{/if} {#if debugEnabled && instanceConnections.length > 0}
{#each instanceConnections as conn}
{conn.from} -> {conn.to}: {conn.ip} ({conn.ifaceLabel})
{/each}
{/if} {#if downloadInfo.isDownloading && downloadInfo.progress}
{downloadInfo.progress.percentage.toFixed( 1, )}% {formatBytes( downloadInfo.progress.downloadedBytes, )}/{formatBytes( downloadInfo.progress.totalBytes, )}
{formatSpeed( downloadInfo.progress.speed, )} ETA: {formatEta( downloadInfo.progress.etaMs, )} {downloadInfo.progress .completedFiles}/{downloadInfo.progress .totalFiles} files
{#if downloadInfo.perNode.length > 0}
{#each downloadInfo.perNode.filter((n) => n.status === "downloading" && n.progress) as nodeProg} {@const nodePercent = Math.min( 100, Math.max(0, nodeProg.percentage), )} {@const isExpanded = instanceDownloadExpandedNodes.has( nodeProg.nodeId, )}
{#if isExpanded}
{#if nodeProg.progress?.files ?? [].length === 0}
No file details reported.
{:else} {#each nodeProg.progress?.files ?? [] as f} {@const filePercent = Math.min( 100, Math.max(0, f.percentage ?? 0), )} {@const isFileComplete = filePercent >= 100}
{f.name} {filePercent.toFixed(1)}%
{formatBytes( f.downloadedBytes, )} / {formatBytes( f.totalBytes, )} {formatSpeed(f.speed)} • ETA {formatEta( f.etaMs, )}
{/each} {/if}
{/if}
{/each}
{/if}
DOWNLOADING

Downloading model files. Model runs on your devices so needs to be downloaded before you can chat.

{:else}
{downloadInfo.statusText}
{#if isLoading} {@const loadStatus = deriveInstanceStatus(instance)} {#if loadStatus.totalLayers && loadStatus.totalLayers > 0}
{( ((loadStatus.layersLoaded ?? 0) / loadStatus.totalLayers) * 100 ).toFixed(0)}% {loadStatus.layersLoaded ?? 0} / {loadStatus.totalLayers} layers
{:else}

Loading model into memory...

{/if} {:else if isWarmingUp}

Warming up...

{:else if isReady || isRunning}

Ready to chat!

{/if}
{#if downloadInfo.isFailed && downloadInfo.errorMessage}
{downloadInfo.errorMessage}
{/if} {/if}
{/each}
{/if}

Load Model

{models.length} models
{#if showAdvancedOptions}
Sharding Strategy:
Interconnect:
Minimum Devices:
{#if availableMinNodes > 1}
{/if} {#each Array.from({ length: availableMinNodes }, (_, i) => i + 1) as n} {@const isValid = validMinNodeCounts().has(n)} {@const isSelected = selectedMinNodes === n} {@const position = availableMinNodes > 1 ? ((n - 1) / (availableMinNodes - 1)) * 100 : 50}
{n}
{/each}
{/if}
{#if models.length === 0}
Loading models...
{:else if loadingPreviews}
Loading preview...
{:else} {@const selectedModel = models.find( (m) => m.id === selectedModelId, )} {@const allPreviews = filteredPreviews()} {#if selectedModel && allPreviews.length > 0} {@const tags = modelTags()[selectedModel.id] || []}
{#each allPreviews as apiPreview, i} {@const downloadStatus = getModelDownloadStatus( selectedModel.id, apiPreview.memory_delta_by_node ? Object.keys(apiPreview.memory_delta_by_node) : undefined, )}
{ if (apiPreview.memory_delta_by_node) { hoveredPreviewNodes = new Set( Object.entries(apiPreview.memory_delta_by_node) .filter(([, delta]) => (delta ?? 0) > 0) .map(([nodeId]) => nodeId), ); } }} onmouseleave={() => (hoveredPreviewNodes = new Set())} > launchInstance(selectedModel.id, apiPreview)} {tags} {apiPreview} modelIdOverride={apiPreview.model_id} />
{/each}
{:else if selectedModel}
No valid configurations for current settings
{/if} {/if}
{/snippet}
{:else}
{#if chatLaunchState !== "idle" && chatLaunchState !== "ready"}
{#if pendingChatModelId}

{pendingChatModelId.split("/").pop()?.replace(/-/g, " ") || pendingChatModelId}

{/if} {#if chatLaunchState === "launching"}

Preparing to launch…

{:else if chatLaunchState === "downloading"}
Downloading {#if chatLaunchDownload} {chatLaunchDownload.percentage.toFixed(1)}% {/if}
{#if chatLaunchDownload}
{formatBytes(chatLaunchDownload.downloadedBytes)} / {formatBytes( chatLaunchDownload.totalBytes, )} {#if chatLaunchDownload.speed > 0} {formatBytes(chatLaunchDownload.speed)}/s {/if} {#if chatLaunchDownload.etaMs > 0} · {formatEta(chatLaunchDownload.etaMs)} {/if}
{/if}
{:else if chatLaunchState === "loading"}
Loading model {#if chatLaunchLoadProgress} {chatLaunchLoadProgress.layersLoaded}/{chatLaunchLoadProgress.totalLayers} layers {/if}
{/if}
{:else if messages().length > 0 || chatLaunchState === "ready"}
{#if chatLaunchState === "ready" && selectedChatCategory} {@const prompts = categorySuggestedPrompts[selectedChatCategory] ?? categorySuggestedPrompts.auto}

Try asking

{#each prompts as prompt} {/each}
{/if}
{:else}
({ id: m.id, name: m.name ?? "", base_model: m.base_model ?? "", storage_size_megabytes: m.storage_size_megabytes ?? 0, capabilities: m.capabilities ?? [], family: m.family ?? "", quantization: m.quantization ?? "", }))} clusterLabel={chatClusterLabel} totalMemoryGB={availableMemoryGB()} onSelect={handleChatModelSelect} onAddModel={handleChatAddModel} />
{/if}
{#if minimized} {/if}
{/if}
{#if !showOnboarding} m.id))} canModelFit={(modelId) => { const model = models.find((m) => m.id === modelId); return model ? hasEnoughMemory(model) : false; }} getModelFitStatus={(modelId): ModelMemoryFitStatus => { const model = models.find((m) => m.id === modelId); return model ? getModelMemoryFitStatus(model) : "too_large"; }} onSelect={(modelId) => { if (modelPickerContext === "chat") { handleChatPickerSelect(modelId); } else { handleModelPickerSelect(modelId); } }} onClose={() => (isModelPickerOpen = false)} onToggleFavorite={toggleFavorite} onAddModel={addModelFromPicker} onDeleteModel={deleteCustomModel} totalMemoryGB={clusterMemory().total / (1024 * 1024 * 1024)} usedMemoryGB={clusterMemory().used / (1024 * 1024 * 1024)} {downloadsData} topologyNodes={data?.nodes} instanceStatuses={modelInstanceStatuses} /> {/if} ================================================ FILE: dashboard/src/routes/downloads/+page.svelte ================================================

Downloads

Overview of models on each node

Last update: {lastUpdateTs ? new Date(lastUpdateTs).toLocaleTimeString() : "n/a"}
{#if !hasDownloads}
No downloads found. Start a model download to see progress here.
Download keys detected: {downloadKeys.length === 0 ? "none" : downloadKeys.join(", ")}
{:else}
{#each nodeColumns as col} {/each} {#each modelRows as row} {#each nodeColumns as col} {@const cell = row.cells[col.nodeId] ?? { kind: "not_present" as const, }} {/each} {/each}
Model
{col.label}
{#if col.diskAvailable != null}
{formatBytes(col.diskAvailable)} free
{/if}
{row.prettyName ?? row.modelId}
{#if row.prettyName}
{row.modelId}
{/if}
{#if cell.kind === "completed"}
{formatBytes(cell.totalBytes)}
{:else if cell.kind === "downloading"}
{clampPercent(cell.percentage).toFixed(1)}%
{formatSpeed(cell.speed)}
{:else if cell.kind === "pending"}
0 ? `${formatBytes(cell.downloaded)} / ${formatBytes(cell.total)} downloaded (paused)` : "Download pending"} > {#if cell.downloaded > 0 && cell.total > 0} {formatBytes(cell.downloaded)} / {formatBytes( cell.total, )}
{#if row.shardMetadata} {:else} paused {/if} {:else if row.shardMetadata} {:else} ... {/if}
{:else if cell.kind === "failed"}
{#if row.shardMetadata} {/if}
{:else}
-- {#if row.shardMetadata} {/if}
{/if}
{/if}
{#if infoRow}
(infoRow = null)} role="presentation" >
{/if} ================================================ FILE: dashboard/src/routes/traces/+page.svelte ================================================

Traces

{#if selectedIds.size > 0} {/if}
{#if loading}
Loading traces...
{:else if error}
{error}
{:else if traces.length === 0}
No traces found.
Run exo with EXO_TRACING_ENABLED=1 to collect traces.
{:else}
{#each traces as trace} {@const isSelected = selectedIds.has(trace.taskId)}
toggleSelect(trace.taskId)} onkeydown={(e) => { if (e.key === "Enter" || e.key === " ") { e.preventDefault(); toggleSelect(trace.taskId); } }} >
e.stopPropagation()} > {trace.taskId}
{formatDate(trace.createdAt)} • {formatBytes( trace.fileSize, )}
e.stopPropagation()} > View Stats
{/each}
{/if}
================================================ FILE: dashboard/src/routes/traces/[taskId]/+page.svelte ================================================

Trace

{taskId}

All Traces
{#if loading}
Loading trace data...
{:else if error}
{error}
{:else if stats}

Summary

{formatDuration(stats.totalWallTimeUs)}
Total wall time
{#if phases.length > 0}

By Phase (avg per node)

{#each phases as phase} {@const normalizedTotal = phase.totalUs / nodeCount} {@const normalizedStepCount = phase.stepCount / nodeCount}
{phase.name} {formatDuration(normalizedTotal)} ({normalizedStepCount} steps, {formatDuration( normalizedTotal / normalizedStepCount, )}/step)
{#if phase.subcategories.length > 0}
{#each phase.subcategories as subcat} {@const normalizedSubcat = subcat.stats.totalUs / nodeCount} {@const pct = formatPercentage( normalizedSubcat, normalizedTotal, )} {@const perStep = normalizedSubcat / normalizedStepCount}
{subcat.name} {formatDuration(normalizedSubcat)} ({pct}) {formatDuration(perStep)}/step
{/each}
{/if}
{/each}
{/if} {#if sortedRanks.length > 0}

By Rank

{#each sortedRanks as rank} {@const rankStats = stats.byRank[rank]} {@const rankPhases = parsePhases(rankStats.byCategory)}
Rank {rank}
{#each rankPhases as phase}
{phase.name} {formatDuration(phase.totalUs)} ({phase.stepCount}x)
{#if phase.subcategories.length > 0}
{#each phase.subcategories as subcat} {@const pct = formatPercentage( subcat.stats.totalUs, phase.totalUs, )} {@const perStep = subcat.stats.totalUs / phase.stepCount}
{subcat.name} {formatDuration(subcat.stats.totalUs)} ({pct}) {formatDuration(perStep)}/step
{/each}
{/if}
{/each}
{/each}
{/if} {/if}
================================================ FILE: dashboard/svelte.config.js ================================================ import adapter from '@sveltejs/adapter-static'; import { vitePreprocess } from '@sveltejs/vite-plugin-svelte'; /** @type {import('@sveltejs/kit').Config} */ const config = { preprocess: [vitePreprocess()], kit: { paths: { relative: true }, router: { type: 'hash' }, adapter: adapter({ pages: 'build', assets: 'build', fallback: 'index.html', precompress: false, strict: true }), alias: { $lib: 'src/lib', $components: 'src/lib/components' } } }; export default config; ================================================ FILE: dashboard/tsconfig.json ================================================ { "extends": "./.svelte-kit/tsconfig.json", "compilerOptions": { "allowJs": true, "checkJs": true, "esModuleInterop": true, "forceConsistentCasingInFileNames": true, "resolveJsonModule": true, "skipLibCheck": true, "sourceMap": true, "strict": true, "moduleResolution": "bundler" } } ================================================ FILE: dashboard/vite.config.ts ================================================ import tailwindcss from "@tailwindcss/vite"; import { sveltekit } from "@sveltejs/kit/vite"; import { defineConfig } from "vite"; export default defineConfig({ plugins: [tailwindcss(), sveltekit()], server: { proxy: { "/v1": "http://localhost:52415", "/state": "http://localhost:52415", "/models": "http://localhost:52415", "/instance": "http://localhost:52415", }, }, }); ================================================ FILE: docs/api.md ================================================ # EXO API – Technical Reference This document describes the REST API exposed by the **EXO** service, as implemented in: `src/exo/master/api.py` The API is used to manage model instances in the cluster, inspect cluster state, and perform inference using multiple API-compatible interfaces. Base URL example: ``` http://localhost:52415 ``` ## 1. General / Meta Endpoints ### Get Master Node ID **GET** `/node_id` Returns the identifier of the current master node. **Response (example):** ```json { "node_id": "node-1234" } ``` ### Get Cluster State **GET** `/state` Returns the current state of the cluster, including nodes and active instances. **Response:** JSON object describing topology, nodes, and instances. ### Get Events **GET** `/events` Returns the list of internal events recorded by the master (mainly for debugging and observability). **Response:** Array of event objects. ## 2. Model Instance Management ### Create Instance **POST** `/instance` Creates a new model instance in the cluster. **Request body (example):** ```json { "instance": { "model_id": "llama-3.2-1b", "placement": { } } } ``` **Response:** JSON description of the created instance. ### Delete Instance **DELETE** `/instance/{instance_id}` Deletes an existing instance by ID. **Path parameters:** * `instance_id`: string, ID of the instance to delete **Response:** Status / confirmation JSON. ### Get Instance **GET** `/instance/{instance_id}` Returns details of a specific instance. **Path parameters:** * `instance_id`: string **Response:** JSON description of the instance. ### Preview Placements **GET** `/instance/previews?model_id=...` Returns possible placement previews for a given model. **Query parameters:** * `model_id`: string, required **Response:** Array of placement preview objects. ### Compute Placement **GET** `/instance/placement` Computes a placement for a potential instance without creating it. **Query parameters (typical):** * `model_id`: string * `sharding`: string or config * `instance_meta`: JSON-encoded metadata * `min_nodes`: integer **Response:** JSON object describing the proposed placement / instance configuration. ### Place Instance (Dry Operation) **POST** `/place_instance` Performs a placement operation for an instance (planning step), without necessarily creating it. **Request body:** JSON describing the instance to be placed. **Response:** Placement result. ## 3. Models ### List Models **GET** `/models` **GET** `/v1/models` (alias) Returns the list of available models and their metadata. **Query parameters:** * `status`: string (optional) - Filter by `downloaded` to show only downloaded models **Response:** Array of model descriptors including `is_custom` field for custom HuggingFace models. ### Add Custom Model **POST** `/models/add` Add a custom model from HuggingFace hub. **Request body (example):** ```json { "model_id": "mlx-community/my-custom-model" } ``` **Response:** Model descriptor for the added model. **Security note:** Models with `trust_remote_code` enabled in their configuration require explicit opt-in (default is false) for security. ### Delete Custom Model **DELETE** `/models/custom/{model_id}` Delete a user-added custom model card. **Path parameters:** * `model_id`: string, ID of the custom model to delete **Response:** Confirmation JSON with deleted model ID. ### Search Models **GET** `/models/search` Search HuggingFace Hub for mlx-community models. **Query parameters:** * `query`: string (optional) - Search query * `limit`: integer (default: 20) - Maximum number of results **Response:** Array of HuggingFace model search results. ## 4. Inference / Chat Completions ### OpenAI-Compatible Chat Completions **POST** `/v1/chat/completions` Executes a chat completion request using an OpenAI-compatible schema. Supports streaming and non-streaming modes. **Request body (example):** ```json { "model": "llama-3.2-1b", "messages": [ { "role": "system", "content": "You are a helpful assistant." }, { "role": "user", "content": "Hello" } ], "stream": false } ``` **Request parameters:** * `model`: string, required - Model ID to use * `messages`: array, required - Conversation messages * `stream`: boolean (default: false) - Enable streaming responses * `max_tokens`: integer (optional) - Maximum tokens to generate * `temperature`: float (optional) - Sampling temperature * `top_p`: float (optional) - Nucleus sampling parameter * `top_k`: integer (optional) - Top-k sampling parameter * `stop`: string or array (optional) - Stop sequences * `seed`: integer (optional) - Random seed for reproducibility * `enable_thinking`: boolean (optional) - Enable thinking mode for capable models (DeepSeek V3.1, Qwen3, GLM-4.7) * `tools`: array (optional) - Tool definitions for function calling * `logprobs`: boolean (optional) - Return log probabilities * `top_logprobs`: integer (optional) - Number of top log probabilities to return **Response:** OpenAI-compatible chat completion response. **Streaming response format:** When `stream=true`, returns Server-Sent Events (SSE) with format: ``` data: {"id":"...","object":"chat.completion","created":...,"model":"...","choices":[...]} data: [DONE] ``` **Non-streaming response includes usage statistics:** ```json { "id": "...", "object": "chat.completion", "created": 1234567890, "model": "llama-3.2-1b", "choices": [{ "index": 0, "message": { "role": "assistant", "content": "Hello! How can I help you?" }, "finish_reason": "stop" }], "usage": { "prompt_tokens": 15, "completion_tokens": 8, "total_tokens": 23 } } ``` **Cancellation:** You can cancel an active generation by closing the HTTP connection. The server detects the disconnection and stops processing. ### Claude Messages API **POST** `/v1/messages` Executes a chat completion request using the Claude Messages API format. Supports streaming and non-streaming modes. **Request body (example):** ```json { "model": "llama-3.2-1b", "messages": [ { "role": "user", "content": "Hello" } ], "max_tokens": 1024, "stream": false } ``` **Streaming response format:** When `stream=true`, returns Server-Sent Events with Claude-specific event types: * `message_start` - Message generation started * `content_block_start` - Content block started * `content_block_delta` - Incremental content chunk * `content_block_stop` - Content block completed * `message_delta` - Message metadata updates * `message_stop` - Message generation completed **Response:** Claude-compatible messages response. ### OpenAI Responses API **POST** `/v1/responses` Executes a chat completion request using the OpenAI Responses API format. Supports streaming and non-streaming modes. **Request body (example):** ```json { "model": "llama-3.2-1b", "messages": [ { "role": "user", "content": "Hello" } ], "stream": false } ``` **Streaming response format:** When `stream=true`, returns Server-Sent Events with response-specific event types: * `response.created` - Response generation started * `response.in_progress` - Response is being generated * `response.output_item.added` - New output item added * `response.output_item.done` - Output item completed * `response.done` - Response generation completed **Response:** OpenAI Responses API-compatible response. ### Benchmarked Chat Completions **POST** `/bench/chat/completions` Same as `/v1/chat/completions`, but also returns performance and generation statistics. **Request body:** Same schema as `/v1/chat/completions`. **Response:** Chat completion plus benchmarking metrics including: * `prompt_tps` - Tokens per second during prompt processing * `generation_tps` - Tokens per second during generation * `prompt_tokens` - Number of prompt tokens * `generation_tokens` - Number of generated tokens * `peak_memory_usage` - Peak memory used during generation ### Cancel Command **POST** `/v1/cancel/{command_id}` Cancels an active generation command (text or image). Notifies workers and closes the stream. **Path parameters:** * `command_id`: string, ID of the command to cancel **Response (example):** ```json { "message": "Command cancelled.", "command_id": "cmd-abc-123" } ``` Returns 404 if the command is not found or already completed. ## 5. Ollama API Compatibility EXO provides Ollama API compatibility for tools like OpenWebUI. ### Ollama Chat **POST** `/ollama/api/chat` **POST** `/ollama/api/api/chat` (alias) **POST** `/ollama/api/v1/chat` (alias) Execute a chat request using Ollama API format. **Request body (example):** ```json { "model": "llama-3.2-1b", "messages": [ { "role": "user", "content": "Hello" } ], "stream": false } ``` **Response:** Ollama-compatible chat response. ### Ollama Generate **POST** `/ollama/api/generate` Execute a text generation request using Ollama API format. **Request body (example):** ```json { "model": "llama-3.2-1b", "prompt": "Hello", "stream": false } ``` **Response:** Ollama-compatible generation response. ### Ollama Tags **GET** `/ollama/api/tags` **GET** `/ollama/api/api/tags` (alias) **GET** `/ollama/api/v1/tags` (alias) Returns list of downloaded models in Ollama tags format. **Response:** Array of model tags with metadata. ### Ollama Show **POST** `/ollama/api/show` Returns model information in Ollama show format. **Request body:** ```json { "name": "llama-3.2-1b" } ``` **Response:** Model details including modelfile and family. ### Ollama PS **GET** `/ollama/api/ps` Returns list of running models (active instances). **Response:** Array of active model instances. ### Ollama Version **GET** `/ollama/api/version` **HEAD** `/ollama/` (alias) **HEAD** `/ollama/api/version` (alias) Returns version information for Ollama API compatibility. **Response:** ```json { "version": "exo v1.0" } ``` ## 6. Image Generation & Editing ### Image Generation **POST** `/v1/images/generations` Executes an image generation request using an OpenAI-compatible schema with additional advanced_params. Supports both streaming and non-streaming modes. **Request body (example):** ```json { "prompt": "a robot playing chess", "model": "exolabs/FLUX.1-dev", "n": 1, "size": "1024x1024", "stream": false, "response_format": "b64_json" } ``` **Request parameters:** * `prompt`: string, required - Text description of the image * `model`: string, required - Image model ID * `n`: integer (default: 1) - Number of images to generate * `size`: string (default: "auto") - Image dimensions. Supported sizes: - `512x512` - `768x768` - `1024x768` - `768x1024` - `1024x1024` - `1024x1536` - `1536x1024` - `1024x1365` - `1365x1024` * `stream`: boolean (default: false) - Enable streaming for partial images * `partial_images`: integer (default: 0) - Number of partial images to stream during generation * `response_format`: string (default: "b64_json") - Either `url` or `b64_json` * `quality`: string (default: "medium") - Either `high`, `medium`, or `low` * `output_format`: string (default: "png") - Either `png`, `jpeg`, or `webp` * `advanced_params`: object (optional) - Advanced generation parameters **Advanced Parameters (`advanced_params`):** | Parameter | Type | Constraints | Description | |-----------|------|-------------|-------------| | `seed` | int | >= 0 | Random seed for reproducible generation | | `num_inference_steps` | int | 1-100 | Number of denoising steps | | `guidance` | float | 1.0-20.0 | Classifier-free guidance scale | | `negative_prompt` | string | - | Text describing what to avoid in the image | **Non-streaming response:** ```json { "created": 1234567890, "data": [ { "b64_json": "iVBORw0KGgoAAAANSUhEUgAA...", "url": null } ] } ``` **Streaming response format:** When `stream=true` and `partial_images > 0`, returns Server-Sent Events: ``` data: {"type":"partial","image_index":0,"partial_index":1,"total_partials":5,"format":"png","data":{"b64_json":"..."}} data: {"type":"final","image_index":0,"format":"png","data":{"b64_json":"..."}} data: [DONE] ``` ### Image Editing **POST** `/v1/images/edits` Executes an image editing request (img2img) using FLUX.1-Kontext-dev or similar models. **Request (multipart/form-data):** * `image`: file, required - Input image to edit * `prompt`: string, required - Text description of desired changes * `model`: string, required - Image editing model ID (e.g., `exolabs/FLUX.1-Kontext-dev`) * `n`: integer (default: 1) - Number of edited images to generate * `size`: string (optional) - Output image dimensions * `response_format`: string (default: "b64_json") - Either `url` or `b64_json` * `input_fidelity`: string (default: "low") - Either `low` or `high` - Controls how closely the output follows the input image * `stream`: string (default: "false") - Enable streaming * `partial_images`: string (default: "0") - Number of partial images to stream * `quality`: string (default: "medium") - Either `high`, `medium`, or `low` * `output_format`: string (default: "png") - Either `png`, `jpeg`, or `webp` * `advanced_params`: string (optional) - JSON-encoded advanced parameters **Response:** Same format as `/v1/images/generations`. ### Benchmarked Image Generation **POST** `/bench/images/generations` Same as `/v1/images/generations`, but also returns generation statistics. **Request body:** Same schema as `/v1/images/generations`. **Response:** Image generation plus benchmarking metrics including: * `seconds_per_step` - Average time per denoising step * `total_generation_time` - Total generation time * `num_inference_steps` - Number of inference steps used * `num_images` - Number of images generated * `image_width` - Output image width * `image_height` - Output image height * `peak_memory_usage` - Peak memory used during generation ### Benchmarked Image Editing **POST** `/bench/images/edits` Same as `/v1/images/edits`, but also returns generation statistics. **Request:** Same schema as `/v1/images/edits`. **Response:** Same format as `/bench/images/generations`, including `generation_stats`. ### List Images **GET** `/images` List all stored images. **Response:** Array of image metadata including URLs and expiration times. ### Get Image **GET** `/images/{image_id}` Retrieve a stored image by ID. **Path parameters:** * `image_id`: string, ID of the image **Response:** Image file with appropriate content type. ## 7. Complete Endpoint Summary ``` # General GET /node_id GET /state GET /events # Instance Management POST /instance GET /instance/{instance_id} DELETE /instance/{instance_id} GET /instance/previews GET /instance/placement POST /place_instance # Models GET /models GET /v1/models POST /models/add DELETE /models/custom/{model_id} GET /models/search # Text Generation (OpenAI Chat Completions) POST /v1/chat/completions POST /bench/chat/completions # Text Generation (Claude Messages API) POST /v1/messages # Text Generation (OpenAI Responses API) POST /v1/responses # Text Generation (Ollama API) POST /ollama/api/chat POST /ollama/api/api/chat POST /ollama/api/v1/chat POST /ollama/api/generate GET /ollama/api/tags GET /ollama/api/api/tags GET /ollama/api/v1/tags POST /ollama/api/show GET /ollama/api/ps GET /ollama/api/version HEAD /ollama/ HEAD /ollama/api/version # Command Control POST /v1/cancel/{command_id} # Image Generation POST /v1/images/generations POST /bench/images/generations POST /v1/images/edits POST /bench/images/edits GET /images GET /images/{image_id} ``` ## 8. Notes ### API Compatibility EXO provides multiple API-compatible interfaces: * **OpenAI Chat Completions API** - Compatible with OpenAI clients and tools * **Claude Messages API** - Compatible with Anthropic's Claude API format * **OpenAI Responses API** - Compatible with OpenAI's Responses API format * **Ollama API** - Compatible with Ollama and tools like OpenWebUI Existing OpenAI, Claude, or Ollama clients can be pointed to EXO by changing the base URL. ### Custom Models You can add custom models from HuggingFace using the `/models/add` endpoint. Custom models are identified by the `is_custom` field in model list responses. **Security:** Models requiring `trust_remote_code` must be explicitly enabled (default is false) for security. Only enable this if you trust the model's remote code. ### Usage Statistics Chat completion responses include usage statistics with: * `prompt_tokens` - Number of tokens in the prompt * `completion_tokens` - Number of tokens generated * `total_tokens` - Sum of prompt and completion tokens ### Request Cancellation You can cancel active requests by: 1. Closing the HTTP connection (for streaming requests) 2. Calling `/v1/cancel/{command_id}` (for any request) The server detects cancellation and stops processing immediately. ### Instance Placement The instance placement endpoints allow you to plan and preview cluster allocations before creating instances. This helps optimize resource usage across nodes. ### Observability The `/events` and `/state` endpoints are primarily intended for operational visibility and debugging. ================================================ FILE: docs/architecture.md ================================================ # EXO Architecture overview EXO uses an _Event Sourcing_ architecture, and Erlang-style _message passing_. To facilitate this, we've written a channel library extending anyio channels with inspiration from tokio::sync::mpsc. Each logical module - designed to be functional independently of the others - communicates with the rest of the system by sending messages on topics. ## Systems There are currently 5 major systems: - Master Executes placement and orders events through a single writer - Worker Schedules work on a node, gathers system information, etc.# - Runner Executes inference jobs (for now) in an isolated process from the worker for fault-tolerance. - API Runs a python webserver for exposing state and commands to client applications - Election Implements a distributed algorithm for master election in unstable networking conditions ## API Layer The API system uses multiple adapters to support multiple API formats, converting them to a single request / response type. ### Adapter Pattern Adapters convert between external API formats and EXO's internal types: ``` Chat Completions → [adapter] → TextGenerationTaskParams → Application Claude Messages → [adapter] → TextGenerationTaskParams → Application Responses API → [adapter] → TextGenerationTaskParams → Application Ollama API → [adapter] → TextGenerationTaskParams → Application ``` Each adapter implements two key functions: 1. **Request conversion**: Converts API-specific requests to `TextGenerationTaskParams` 2. **Response generation**: Converts internal `TokenChunk` streams back to API-specific formats (streaming and non-streaming) ## Topics There are currently 5 topics: - Commands The API and Worker instruct the master when the event log isn't sufficient. Namely placement and catchup requests go through Commands atm. - Local Events All nodes write events here, the master reads those events and orders them - Global Events The master writes events here, all nodes read from this topic and fold the produced events into their `State` - Election Messages Before establishing a cluster, nodes communicate here to negotiate a master node. - Connection Messages The networking system write mdns-discovered hardware connections here. ## Event Sourcing Lots has been written about event sourcing, but it lets us centralize faulty connections and message ACKing with the following model. Whenever a device produces side effects, it captures those side effects in an `Event`. `Event`s are then "applied" to their model of `State`, which is globally distributed across the cluster. Whenever a command is received, it is combined with state to produce side effects, captured in yet more events. The rule of thumb is "`Event`s are past tense, `Command`s are imperative". Telling a node to perform some action like "place this model" or "Give me a copy of the event log" is represented by a command (The worker's `Task`s are also commands), while "this node is using 300GB of ram" is an event. Notably, `Event`s SHOULD never cause side effects on their own. There are a few exceptions to this, we're working out the specifics of generalizing the distributed event sourcing model to make it better suit our needs ## Purity A significant goal of the current design is to make data flow explicit. Classes should either represent simple data (`CamelCaseModel`s typically, and `TaggedModel`s for unions) or active `System`s (Erlang `Actor`s), with all transformations of that data being "referentially transparent" - destructure and construct new data, don't mutate in place. We have had varying degrees of success with this, and are still exploring where purity makes sense. ================================================ FILE: flake.nix ================================================ { description = "The development environment for Exo"; inputs = { nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable"; flake-parts = { url = "github:hercules-ci/flake-parts"; inputs.nixpkgs-lib.follows = "nixpkgs"; }; crane.url = "github:ipetkov/crane"; fenix = { url = "github:nix-community/fenix"; inputs.nixpkgs.follows = "nixpkgs"; }; treefmt-nix = { url = "github:numtide/treefmt-nix"; inputs.nixpkgs.follows = "nixpkgs"; }; dream2nix = { url = "github:nix-community/dream2nix"; inputs.nixpkgs.follows = "nixpkgs"; inputs.pyproject-nix.follows = "pyproject-nix"; }; # Python packaging with uv2nix pyproject-nix = { url = "github:pyproject-nix/pyproject.nix"; inputs.nixpkgs.follows = "nixpkgs"; }; uv2nix = { url = "github:pyproject-nix/uv2nix"; inputs.pyproject-nix.follows = "pyproject-nix"; inputs.nixpkgs.follows = "nixpkgs"; }; pyproject-build-systems = { url = "github:pyproject-nix/build-system-pkgs"; inputs.pyproject-nix.follows = "pyproject-nix"; inputs.uv2nix.follows = "uv2nix"; inputs.nixpkgs.follows = "nixpkgs"; }; # Pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs) nixpkgs-swift.url = "github:NixOS/nixpkgs/08dacfca559e1d7da38f3cf05f1f45ee9bfd213c"; }; nixConfig = { extra-trusted-public-keys = "exo.cachix.org-1:okq7hl624TBeAR3kV+g39dUFSiaZgLRkLsFBCuJ2NZI="; extra-substituters = "https://exo.cachix.org"; }; outputs = inputs: inputs.flake-parts.lib.mkFlake { inherit inputs; } { systems = [ "x86_64-linux" "aarch64-darwin" "aarch64-linux" ]; imports = [ inputs.treefmt-nix.flakeModule ./dashboard/parts.nix ./rust/parts.nix ./python/parts.nix ]; perSystem = { config, self', inputs', pkgs, lib, system, ... }: let # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs) pkgsSwift = import inputs.nixpkgs-swift { inherit system; }; in { # Allow unfree for metal-toolchain (needed for Darwin Metal packages) _module.args.pkgs = import inputs.nixpkgs { inherit system; config.allowUnfreePredicate = pkg: (pkg.pname or "") == "metal-toolchain"; overlays = [ (import ./nix/apple-sdk-overlay.nix) ]; }; treefmt = { projectRootFile = "flake.nix"; programs = { nixpkgs-fmt.enable = true; ruff-format = { enable = true; excludes = [ "rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi" ]; }; rustfmt = { enable = true; package = config.rust.toolchain; }; prettier = { enable = true; package = self'.packages.prettier-svelte; includes = [ "*.ts" "*.svelte" ]; }; swift-format = { enable = true; package = pkgsSwift.swiftPackages.swift-format; }; shfmt.enable = true; taplo.enable = true; }; }; packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin ( let uvLock = builtins.fromTOML (builtins.readFile ./uv.lock); mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package); uvLockMlxVersion = mlxPackage.version; uvLockMlxRev = builtins.elemAt (builtins.split "#" mlxPackage.source.git) 2; in { metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { }; mlx = pkgs.callPackage ./nix/mlx.nix { inherit (self'.packages) metal-toolchain; inherit uvLockMlxVersion uvLockMlxRev; }; default = self'.packages.exo; } ); devShells.default = with pkgs; pkgs.mkShell { inputsFrom = [ self'.checks.cargo-build ]; packages = [ # FORMATTING config.treefmt.build.wrapper # PYTHON python313 uv ruff basedpyright # RUST config.rust.toolchain maturin # NIX nixpkgs-fmt # SVELTE nodejs # MISC just jq ] ++ lib.optionals stdenv.isLinux [ unixtools.ifconfig ] ++ lib.optionals stdenv.isDarwin [ macmon ]; OPENSSL_NO_VENDOR = "1"; shellHook = '' export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib" ${lib.optionalString stdenv.isLinux '' export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH" ''} ''; }; }; }; } ================================================ FILE: justfile ================================================ export NIX_CONFIG := "extra-experimental-features = nix-command flakes" fmt: treefmt || nix fmt lint: uv run ruff check --fix test: uv run pytest src check: uv run basedpyright --project pyproject.toml sync: uv sync --all-packages sync-clean: uv sync --all-packages --force-reinstall --no-cache rust-rebuild: cargo run --bin stub_gen uv sync --reinstall-package exo_pyo3_bindings build-dashboard: #!/usr/bin/env bash cd dashboard npm install npm run build package: uv run pyinstaller packaging/pyinstaller/exo.spec clean: rm -rf **/__pycache__ rm -rf target/ rm -rf .venv rm -rf dashboard/node_modules rm -rf dashboard/.svelte-kit rm -rf dashboard/build ================================================ FILE: nix/apple-sdk/metadata/versions.json ================================================ { "14": { "urls": [ "https://swcdn.apple.com/content/downloads/14/48/052-59890-A_I0F5YGAY0Y/p9n40hio7892gou31o1v031ng6fnm9sb3c/CLTools_macOSNMOS_SDK.pkg", "https://web.archive.org/web/20250211001355/https://swcdn.apple.com/content/downloads/14/48/052-59890-A_I0F5YGAY0Y/p9n40hio7892gou31o1v031ng6fnm9sb3c/CLTools_macOSNMOS_SDK.pkg" ], "version": "14.4", "hash": "sha256-QozDiwY0Czc0g45vPD7G4v4Ra+3DujCJbSads3fJjjM=" }, "15": { "urls": [ "https://swcdn.apple.com/content/downloads/52/01/082-41241-A_0747ZN8FHV/dectd075r63pppkkzsb75qk61s0lfee22j/CLTools_macOSNMOS_SDK.pkg", "https://web.archive.org/web/20250530132510/https://swcdn.apple.com/content/downloads/52/01/082-41241-A_0747ZN8FHV/dectd075r63pppkkzsb75qk61s0lfee22j/CLTools_macOSNMOS_SDK.pkg" ], "version": "15.5", "hash": "sha256-HBiSJuw1XBUK5R/8Sj65c3rftSEvQl/O9ZZVp/g1Amo=" }, "26": { "urls": [ "https://swcdn.apple.com/content/downloads/60/22/089-71960-A_W8BL1RUJJ6/5zkyplomhk1cm7z6xja2ktgapnhhti6wwd/CLTools_macOSNMOS_SDK.pkg", "https://web.archive.org/web/20250915230423/https://swcdn.apple.com/content/downloads/60/22/089-71960-A_W8BL1RUJJ6/5zkyplomhk1cm7z6xja2ktgapnhhti6wwd/CLTools_macOSNMOS_SDK.pkg" ], "version": "26.2", "hash": "sha256-hXRlMieVv0smna5uiWRwq87IWOaPWtAjAldbi+wQXcw=" } } ================================================ FILE: nix/apple-sdk-overlay.nix ================================================ # Overlay that builds apple-sdk with a custom versions.json (for SDK 26.2). # The upstream nixpkgs package reads versions.json at eval time via a relative # path, so we can't override it through callPackage args. Instead, we copy # the upstream source and patch the one file. final: _prev: let upstreamSrc = final.path + "/pkgs/by-name/ap/apple-sdk"; patchedSrc = final.runCommandLocal "apple-sdk-src-patched" { } '' cp -r ${upstreamSrc} $out chmod -R u+w $out cp ${./apple-sdk/metadata/versions.json} $out/metadata/versions.json ''; in { apple-sdk_26 = final.callPackage (patchedSrc + "/package.nix") { darwinSdkMajorVersion = "26"; }; } ================================================ FILE: nix/darwin-build-fixes.patch ================================================ diff --git a/CMakeLists.txt b/CMakeLists.txt index 0ed30932..d8528132 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -177,11 +177,7 @@ if(MLX_BUILD_METAL) add_compile_definitions(MLX_METAL_DEBUG) endif() - # Throw an error if xcrun not found - execute_process( - COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version" - OUTPUT_VARIABLE MACOS_SDK_VERSION - OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY) + set(MACOS_SDK_VERSION @sdkVersion@) if(${MACOS_SDK_VERSION} LESS 14.0) message( @@ -199,11 +195,8 @@ if(MLX_BUILD_METAL) endif() set(XCRUN_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}") endif() - execute_process( - COMMAND - zsh "-c" - "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'" - OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY) + set( + MLX_METAL_VERSION @metalVersion@) FetchContent_Declare(metal_cpp URL ${METAL_CPP_URL}) FetchContent_MakeAvailable(metal_cpp) target_include_directories( diff --git a/cmake/extension.cmake b/cmake/extension.cmake index 13db804a..5b385132 100644 --- a/cmake/extension.cmake +++ b/cmake/extension.cmake @@ -36,7 +36,7 @@ macro(mlx_build_metallib) add_custom_command( OUTPUT ${MTLLIB_BUILD_TARGET} COMMAND - xcrun -sdk macosx metal + metal -fmodules-cache-path=${CMAKE_BINARY_DIR}/metal-cache "$" ${MTLLIB_COMPILE_OPTIONS} ${MTLLIB_SOURCES} -o ${MTLLIB_BUILD_TARGET} DEPENDS ${MTLLIB_DEPS} ${MTLLIB_SOURCES} diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt index 262b0495..5c7446ad 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/mlx/backend/metal/kernels/CMakeLists.txt @@ -29,7 +29,7 @@ function(build_kernel_base TARGET SRCFILE DEPS) "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}") endif() add_custom_command( - COMMAND xcrun -sdk macosx metal ${METAL_FLAGS} -c ${SRCFILE} + COMMAND metal -fmodules-cache-path=${CMAKE_BINARY_DIR}/metal-cache ${METAL_FLAGS} -c ${SRCFILE} -I${PROJECT_SOURCE_DIR} -o ${TARGET}.air DEPENDS ${SRCFILE} ${DEPS} ${BASE_HEADERS} OUTPUT ${TARGET}.air @@ -170,7 +170,7 @@ endif() add_custom_command( OUTPUT ${MLX_METAL_PATH}/mlx.metallib - COMMAND xcrun -sdk macosx metallib ${KERNEL_AIR} -o + COMMAND metallib ${KERNEL_AIR} -o ${MLX_METAL_PATH}/mlx.metallib DEPENDS ${KERNEL_AIR} COMMENT "Building mlx.metallib" diff --git a/mlx/backend/metal/make_compiled_preamble.sh b/mlx/backend/metal/make_compiled_preamble.sh index bb55ed3a..94ea7dd7 100644 --- a/mlx/backend/metal/make_compiled_preamble.sh +++ b/mlx/backend/metal/make_compiled_preamble.sh @@ -31,7 +31,7 @@ OUTPUT_FILE=${OUTPUT_DIR}/${SRC_NAME}.cpp mkdir -p "$OUTPUT_DIR" # Use the metal compiler to get a list of headers (with depth) -CCC="xcrun -sdk macosx metal -x metal" +CCC="metal -x metal -fmodules-cache-path=${OUTPUT_DIR}/metal-cache" HDRS=$( $CCC -I"$SRC_DIR" -I"$JIT_INCLUDES" -DMLX_METAL_JIT -E -P -CC -C -H "$INPUT_FILE" $CFLAGS -w 2>&1 1>/dev/null ) # Remove any included system frameworks (for MetalPerformancePrimitive headers) ================================================ FILE: nix/metal-toolchain.nix ================================================ { lib, stdenvNoCC, requireFile, nix }: let narFile = requireFile { name = "metal-toolchain-17C48.nar"; message = '' The Metal Toolchain NAR must be available. If you have cachix configured for exo.cachix.org, this should be automatic. Otherwise: 1. Install Xcode 26+ from the App Store 2. Run: xcodebuild -downloadComponent MetalToolchain 3. Export the toolchain: hdiutil attach "$(find /System/Library/AssetsV2/com_apple_MobileAsset_MetalToolchain -name '*.dmg' | head -1)" -mountpoint /tmp/metal-dmg cp -R /tmp/metal-dmg/Metal.xctoolchain /tmp/metal-export hdiutil detach /tmp/metal-dmg 4. Create NAR and add to store: nix nar pack /tmp/metal-export > /tmp/metal-toolchain-17C48.nar nix store add --mode flat /tmp/metal-toolchain-17C48.nar ''; hash = "sha256-ayR5mXN4sZAddwKEG2OszGRF93k9ZFc7H0yi2xbylQw="; }; in stdenvNoCC.mkDerivation { pname = "metal-toolchain"; version = "17C48"; dontUnpack = true; dontBuild = true; dontFixup = true; nativeBuildInputs = [ nix ]; installPhase = '' runHook preInstall nix-store --restore $out < ${narFile} # Create bin directory with symlinks for PATH mkdir -p $out/bin ln -s $out/usr/bin/metal $out/bin/metal ln -s $out/usr/bin/metallib $out/bin/metallib runHook postInstall ''; # Metal language version for CMake (from: echo __METAL_VERSION__ | metal -E -x metal -P -) passthru.metalVersion = "400"; meta = { description = "Apple Metal compiler toolchain"; platforms = [ "aarch64-darwin" ]; license = lib.licenses.unfree; }; } ================================================ FILE: nix/mlx.nix ================================================ { stdenv , lib , fetchFromGitHub , replaceVars , fetchzip , cmake , nlohmann_json , apple-sdk_26 , metal-toolchain , runCommand , fmt , python313Packages , uvLockMlxVersion , uvLockMlxRev }: assert stdenv.isDarwin; let python = python313Packages.python; # Static dependencies included directly during compilation gguf-tools = fetchFromGitHub { owner = "antirez"; repo = "gguf-tools"; rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848"; hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA="; }; metal_cpp = fetchzip { url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip"; hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8="; }; nanobind = fetchFromGitHub { owner = "wjakob"; repo = "nanobind"; rev = "v2.10.2"; hash = "sha256-io44YhN+VpfHFWyvvLWSanRgbzA0whK8WlDNRi3hahU="; fetchSubmodules = true; }; mlx = stdenv.mkDerivation rec { pname = "mlx"; version = uvLockMlxVersion; pyproject = true; src = fetchFromGitHub { owner = "rltakashige"; repo = "mlx-jaccl-fix-small-recv"; rev = uvLockMlxRev; hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ="; }; patches = [ (replaceVars ./darwin-build-fixes.patch { sdkVersion = apple-sdk_26.version; metalVersion = metal-toolchain.metalVersion; }) ]; postPatch = '' substituteInPlace mlx/backend/cpu/jit_compiler.cpp \ --replace-fail "g++" "$CXX" ''; dontUseCmakeConfigure = true; enableParallelBuilding = true; # Allows multiple cores to be used in Python builds. postUnpack = '' export MAKEFLAGS+="''${enableParallelBuilding:+-j$NIX_BUILD_CORES}" ''; # Updates the wrong fetcher rev attribute passthru.skipBulkUpdate = true; env = { DEV_RELEASE = 1; CMAKE_ARGS = toString [ (lib.cmakeBool "USE_SYSTEM_FMT" true) (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}") (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json.src}") (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}") (lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true) (lib.cmakeBool "MLX_BUILD_CPU" true) (lib.cmakeBool "MLX_BUILD_METAL" true) (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}") (lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${apple-sdk_26.version}") (lib.cmakeOptionType "filepath" "CMAKE_OSX_SYSROOT" "${apple-sdk_26.passthru.sdkroot}") ]; SDKROOT = apple-sdk_26.passthru.sdkroot; MACOSX_DEPLOYMENT_TARGET = apple-sdk_26.version; }; build-system = [ python313Packages.setuptools ]; nativeBuildInputs = [ cmake metal-toolchain python313Packages.pypaBuildHook python313Packages.pypaInstallHook python313Packages.setuptools python313Packages.typing-extensions python313Packages.wheel python313Packages.cmake python313Packages.ninja ]; buildInputs = [ fmt gguf-tools python313Packages.nanobind python313Packages.pybind11 apple-sdk_26 ]; # Tests require Metal GPU access which isn't available in the Nix sandbox. # To run tests, build with: nix build --option sandbox false .#mlx.passthru.tests.mlxTest doCheck = false; pythonImportsCheck = [ "mlx" ]; passthru.tests = { # Runs example scripts to verify MLX works. Requires --option sandbox false # since Metal GPU access is needed. mlxTest = runCommand "run-mlx-examples" { buildInputs = [ mlx ]; nativeBuildInputs = [ python ]; } '' cp ${src}/examples/python/logistic_regression.py . ${python.interpreter} logistic_regression.py rm logistic_regression.py cp ${src}/examples/python/linear_regression.py . ${python.interpreter} linear_regression.py rm linear_regression.py touch $out ''; }; meta = { homepage = "https://github.com/ml-explore/mlx"; description = "Array framework for Apple silicon"; changelog = "https://github.com/ml-explore/mlx/releases/tag/${src.tag}"; license = lib.licenses.mit; platforms = [ "aarch64-darwin" ]; }; }; in mlx ================================================ FILE: packaging/dmg/create-dmg.sh ================================================ #!/usr/bin/env bash # create-dmg.sh — Build a polished macOS DMG installer for EXO # # Usage: # ./packaging/dmg/create-dmg.sh [volume-name] # # Example: # ./packaging/dmg/create-dmg.sh output/EXO.app EXO-1.0.0.dmg "EXO" # # Creates a DMG with: # - Custom background image with drag-to-Applications arrow # - App icon on left, Applications alias on right # - Proper window size and icon positioning set -euo pipefail APP_PATH="${1:?Usage: create-dmg.sh [volume-name]}" OUTPUT_DMG="${2:?Usage: create-dmg.sh [volume-name]}" VOLUME_NAME="${3:-EXO}" SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" BACKGROUND_SCRIPT="${SCRIPT_DIR}/generate-background.py" TEMP_DIR="$(mktemp -d)" DMG_STAGING="${TEMP_DIR}/dmg-root" TEMP_DMG="${TEMP_DIR}/temp.dmg" BACKGROUND_PNG="${TEMP_DIR}/background.png" cleanup() { rm -rf "$TEMP_DIR"; } trap cleanup EXIT echo "==> Creating DMG installer for ${VOLUME_NAME}" # ── Step 1: Generate background image ──────────────────────────────────────── if command -v python3 &>/dev/null; then python3 "$BACKGROUND_SCRIPT" "$BACKGROUND_PNG" echo " Background image generated" else echo " Warning: python3 not found, skipping custom background" BACKGROUND_PNG="" fi # ── Step 2: Prepare staging directory ───────────────────────────────────────── mkdir -p "$DMG_STAGING" cp -R "$APP_PATH" "$DMG_STAGING/" ln -s /Applications "$DMG_STAGING/Applications" # ── Step 3: Create writable DMG ────────────────────────────────────────────── # Calculate required size (app size + 20MB headroom) APP_SIZE_KB=$(du -sk "$APP_PATH" | cut -f1) DMG_SIZE_KB=$((APP_SIZE_KB + 20480)) hdiutil create \ -volname "$VOLUME_NAME" \ -size "${DMG_SIZE_KB}k" \ -fs HFS+ \ -layout SPUD \ "$TEMP_DMG" # ── Step 4: Mount and configure ────────────────────────────────────────────── MOUNT_DIR=$(hdiutil attach "$TEMP_DMG" -readwrite -noverify | awk -F'\t' '/Apple_HFS/ {gsub(/^[[:space:]]+|[[:space:]]+$/, "", $NF); print $NF}') echo " Mounted at: $MOUNT_DIR" # Copy contents cp -R "$DMG_STAGING/"* "$MOUNT_DIR/" # Add background image if [[ -n $BACKGROUND_PNG && -f $BACKGROUND_PNG ]]; then mkdir -p "$MOUNT_DIR/.background" cp "$BACKGROUND_PNG" "$MOUNT_DIR/.background/background.png" fi # ── Step 5: Configure window appearance via AppleScript ────────────────────── # Window: 800×400, app icon on left, Applications on right (matches Ollama layout) # Background image is 1600×740 (2× retina for 800×400 logical window). APP_NAME="$(basename "$APP_PATH")" osascript < DMG created: $OUTPUT_DMG" echo " Size: $(du -h "$OUTPUT_DMG" | cut -f1)" ================================================ FILE: packaging/dmg/generate-background.py ================================================ #!/usr/bin/env python3 """Generate the DMG background image with a centered drag-to-Applications arrow. The output is a 1600×740 retina PNG (2× for 800×400 logical window). Icons are positioned at (200, 190) and (600, 190) in logical coordinates; the arrow is drawn centered between them. Usage: python3 generate-background.py [output.png] If no output path is given, overwrites the bundled background.png in-place. """ from __future__ import annotations import math import sys from pathlib import Path from PIL import Image, ImageDraw # Retina dimensions (2× logical 800×400) WIDTH = 1600 HEIGHT = 740 # Icon positions in logical coords → retina coords # App icon at (200, 190), Applications at (600, 190) APP_X = 200 * 2 # 400 APPS_X = 600 * 2 # 1200 ICON_Y = 190 * 2 # 380 # Arrow drawn between icons, slightly above icon center ARROW_START_X = APP_X + 160 # past the icon ARROW_END_X = APPS_X - 160 # before the Applications icon ARROW_Y = ICON_Y # same height as icons ARROW_RISE = 120 # upward arc height def draw_arrow(draw: ImageDraw.ImageDraw) -> None: """Draw a hand-drawn-style curved arrow from app icon toward Applications.""" color = (30, 30, 30) line_width = 8 # Compute bezier curve points for a gentle upward arc points: list[tuple[float, float]] = [] steps = 80 for i in range(steps + 1): t = i / steps # Quadratic bezier: start → control → end cx = (ARROW_START_X + ARROW_END_X) / 2 cy = ARROW_Y - ARROW_RISE x = (1 - t) ** 2 * ARROW_START_X + 2 * (1 - t) * t * cx + t**2 * ARROW_END_X y = (1 - t) ** 2 * ARROW_Y + 2 * (1 - t) * t * cy + t**2 * ARROW_Y points.append((x, y)) # Draw the curve as connected line segments for i in range(len(points) - 1): draw.line([points[i], points[i + 1]], fill=color, width=line_width) # Arrowhead at the end end_x, end_y = points[-1] # Direction from second-to-last to last point prev_x, prev_y = points[-3] angle = math.atan2(end_y - prev_y, end_x - prev_x) head_len = 36 head_angle = math.radians(25) left_x = end_x - head_len * math.cos(angle - head_angle) left_y = end_y - head_len * math.sin(angle - head_angle) right_x = end_x - head_len * math.cos(angle + head_angle) right_y = end_y - head_len * math.sin(angle + head_angle) draw.polygon( [(end_x, end_y), (left_x, left_y), (right_x, right_y)], fill=color, ) def generate_background(output_path: str) -> None: """Generate a white DMG background with a centered arrow.""" img = Image.new("RGBA", (WIDTH, HEIGHT), (255, 255, 255, 255)) draw = ImageDraw.Draw(img) draw_arrow(draw) img.save(output_path, "PNG") if __name__ == "__main__": default_output = str(Path(__file__).parent / "background.png") out = sys.argv[1] if len(sys.argv) >= 2 else default_output generate_background(out) print(f"Background image written to {out}") ================================================ FILE: packaging/pyinstaller/exo.spec ================================================ # -*- mode: python ; coding: utf-8 -*- import importlib.util import shutil from pathlib import Path from PyInstaller.utils.hooks import collect_submodules PROJECT_ROOT = Path.cwd() SOURCE_ROOT = PROJECT_ROOT / "src" ENTRYPOINT = SOURCE_ROOT / "exo" / "__main__.py" DASHBOARD_DIR = PROJECT_ROOT / "dashboard" / "build" RESOURCES_DIR = PROJECT_ROOT / "resources" EXO_SHARED_MODELS_DIR = SOURCE_ROOT / "exo" / "shared" / "models" if not ENTRYPOINT.is_file(): raise SystemExit(f"Unable to locate Exo entrypoint: {ENTRYPOINT}") if not DASHBOARD_DIR.is_dir(): raise SystemExit(f"Dashboard assets are missing: {DASHBOARD_DIR}") if not RESOURCES_DIR.is_dir(): raise SystemExit(f"Resource assets are missing: {RESOURCES_DIR}") if not EXO_SHARED_MODELS_DIR.is_dir(): raise SystemExit(f"Shared model assets are missing: {EXO_SHARED_MODELS_DIR}") block_cipher = None def _module_directory(module_name: str) -> Path: spec = importlib.util.find_spec(module_name) if spec is None: raise SystemExit(f"Module '{module_name}' is not available in the current environment.") if spec.submodule_search_locations: return Path(next(iter(spec.submodule_search_locations))).resolve() if spec.origin: return Path(spec.origin).resolve().parent raise SystemExit(f"Unable to determine installation directory for '{module_name}'.") MLX_PACKAGE_DIR = _module_directory("mlx") MLX_LIB_DIR = MLX_PACKAGE_DIR / "lib" if not MLX_LIB_DIR.is_dir(): raise SystemExit(f"mlx Metal libraries are missing: {MLX_LIB_DIR}") def _safe_collect(package_name: str) -> list[str]: try: return collect_submodules(package_name) except ImportError: return [] HIDDEN_IMPORTS = sorted( set( collect_submodules("mlx") + _safe_collect("mlx_lm") + _safe_collect("transformers") ) ) DATAS: list[tuple[str, str]] = [ (str(DASHBOARD_DIR), "dashboard"), (str(RESOURCES_DIR), "resources"), (str(MLX_LIB_DIR), "mlx/lib"), (str(EXO_SHARED_MODELS_DIR), "exo/shared/models"), ] MACMON_PATH = shutil.which("macmon") if MACMON_PATH is None: raise SystemExit( "macmon binary not found in PATH. " "Install it via: brew install macmon" ) BINARIES: list[tuple[str, str]] = [ (MACMON_PATH, "."), ] a = Analysis( [str(ENTRYPOINT)], pathex=[str(SOURCE_ROOT)], binaries=BINARIES, datas=DATAS, hiddenimports=HIDDEN_IMPORTS, hookspath=[], hooksconfig={}, runtime_hooks=[], excludes=[], win_no_prefer_redirects=False, win_private_assemblies=False, noarchive=False, ) pyz = PYZ(a.pure, a.zipped_data, cipher=block_cipher) exe = EXE( pyz, a.scripts, [], exclude_binaries=True, name="exo", debug=False, bootloader_ignore_signals=False, strip=False, upx=False, console=True, disable_windowed_traceback=False, argv_emulation=False, target_arch=None, codesign_identity=None, entitlements_file=None, ) coll = COLLECT( exe, a.binaries, a.zipfiles, a.datas, strip=False, upx=False, upx_exclude=[], name="exo", ) ================================================ FILE: pyproject.toml ================================================ [project] name = "exo" version = "0.3.68" description = "Exo" readme = "README.md" requires-python = ">=3.13" dependencies = [ "aiofiles>=24.1.0", "aiohttp>=3.12.14", "types-aiofiles>=24.1.0.20250708", "pydantic>=2.11.7", "fastapi>=0.116.1", "filelock>=3.18.0", "rustworkx>=0.17.1", "huggingface-hub>=0.33.4", "psutil>=7.0.0", "loguru>=0.7.3", "exo_pyo3_bindings", # rust bindings "anyio==4.11.0", "mlx; sys_platform == 'darwin'", "mlx[cpu]==0.30.6; sys_platform == 'linux'", "mlx-lm", "tiktoken>=0.12.0", # required for kimi k2 tokenizer "hypercorn>=0.18.0", "openai-harmony>=0.0.8", "httpx>=0.28.1", "tomlkit>=0.14.0", "mflux==0.16.9", "python-multipart>=0.0.21", "msgspec>=0.19.0", "zstandard>=0.23.0", ] [project.scripts] exo = "exo.main:main" # dependencies only required for development [dependency-groups] dev = [ "basedpyright>=1.29.0", "pyinstaller>=6.17.0", "pytest>=8.4.0", "pytest-asyncio>=1.0.0", "pytest-env", "ruff>=0.11.13", ] # mlx[cuda] requires a newer version of mlx. the ideal on linux is: default to mlx[cpu] unless[cuda] specified. [project.optional-dependencies] # cuda = [ # "mlx[cuda]==0.26.3", # ] ### # workspace configuration ### [tool.uv.workspace] members = ["rust/exo_pyo3_bindings", "bench"] [tool.uv.sources] exo_pyo3_bindings = { workspace = true } mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" } mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" } # Uncomment to use local mlx/mlx-lm development versions: # mlx = { path = "/Users/Shared/mlx", editable=true } # mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true } [build-system] requires = ["uv_build>=0.8.9,<0.9.0"] build-backend = "uv_build" ### # type-checker configuration ### [tool.basedpyright] include = [".venv/lib/mlx", ".venv/lib/mlx_lm", "src", "bench"] typeCheckingMode = "strict" failOnWarnings = true reportAny = "error" reportUnknownVariableType = "error" reportUnknownParameterType = "error" reportMissingParameterType = "error" reportMissingTypeStubs = "error" reportInvalidCast = "error" reportUnnecessaryCast = "error" reportUnnecessaryTypeIgnoreComment = "error" pythonVersion = "3.13" pythonPlatform = "Darwin" exclude = [ "**/.venv", "**/venv", "**/__pycache__", "**/exo_scripts", "**/.direnv", "**/rust", "**/.github", ] stubPath = ".mlx_typings" [[tool.basedpyright.executionEnvironments]] root = "src" ### # uv configuration ### # supported platforms for this project [tool.uv] required-version = ">=0.8.6" prerelease = "allow" environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"] ### # ruff configuration ### [tool.ruff] extend-exclude = [ "shared/protobufs/**", "*mlx_typings/**", "rust/exo_pyo3_bindings/**", "bench/vendor/**", ] [tool.ruff.lint] extend-select = ["I", "N", "B", "A", "PIE", "SIM"] [tool.pytest.ini_options] pythonpath = "." asyncio_mode = "auto" markers = ["slow: marks tests as slow (deselected by default)"] env = ["EXO_TESTS=1"] addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py" filterwarnings = ["ignore:builtin type Swig:DeprecationWarning"] ================================================ FILE: python/parts.nix ================================================ { inputs, ... }: { perSystem = { config, self', pkgs, lib, system, ... }: let # Load workspace from uv.lock workspace = inputs.uv2nix.lib.workspace.loadWorkspace { workspaceRoot = inputs.self; }; # Create overlay from workspace # Use wheels from PyPI for most packages; we override mlx with our pure Nix Metal build overlay = workspace.mkPyprojectOverlay { sourcePreference = "wheel"; }; # Override overlay to inject Nix-built components exoOverlay = final: prev: { # Replace workspace exo_pyo3_bindings with Nix-built wheel. # Preserve passthru so mkVirtualEnv can resolve dependency groups. # Copy .pyi stub + py.typed marker so basedpyright can find the types. exo-pyo3-bindings = pkgs.stdenv.mkDerivation { pname = "exo-pyo3-bindings"; version = "0.1.0"; src = self'.packages.exo_pyo3_bindings; # Install from pre-built wheel nativeBuildInputs = [ final.pyprojectWheelHook ]; dontStrip = true; passthru = prev.exo-pyo3-bindings.passthru or { }; postInstall = '' local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/ touch $siteDir/py.typed ''; }; }; python = pkgs.python313; # Overlay to provide build systems and custom packages buildSystemsOverlay = final: prev: { # mlx-lm is a git dependency that needs setuptools mlx-lm = prev.mlx-lm.overrideAttrs (old: { nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ final.setuptools ]; }); # rouge-score and sacrebleu don't declare setuptools as a build dependency rouge-score = prev.rouge-score.overrideAttrs (old: { nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ final.setuptools ]; }); sacrebleu = prev.sacrebleu.overrideAttrs (old: { nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ final.setuptools ]; }); sqlitedict = prev.sqlitedict.overrideAttrs (old: { nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ final.setuptools ]; }); word2number = prev.word2number.overrideAttrs (old: { nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ final.setuptools ]; }); } // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin { # Use our pure Nix-built MLX with Metal support (macOS only) mlx = self'.packages.mlx; }; # Additional overlay for Linux-specific fixes (type checking env). # Native wheels have shared lib dependencies we don't need at type-check time. linuxOverlay = final: prev: let ignoreMissing = drv: drv.overrideAttrs { autoPatchelfIgnoreMissingDeps = [ "*" ]; }; nvidiaPackages = lib.filterAttrs (name: _: lib.hasPrefix "nvidia-" name) prev; in lib.optionalAttrs pkgs.stdenv.hostPlatform.isLinux ( (lib.mapAttrs (_: ignoreMissing) nvidiaPackages) // { mlx = ignoreMissing prev.mlx; mlx-cuda-13 = prev.mlx-cuda-13.overrideAttrs (old: { buildInputs = (old.buildInputs or [ ]) ++ [ final.nvidia-cublas final.nvidia-cuda-nvrtc final.nvidia-cudnn-cu13 final.nvidia-nccl-cu13 ]; preFixup = '' addAutoPatchelfSearchPath ${final.nvidia-cublas} addAutoPatchelfSearchPath ${final.nvidia-cuda-nvrtc} addAutoPatchelfSearchPath ${final.nvidia-cudnn-cu13} addAutoPatchelfSearchPath ${final.nvidia-nccl-cu13} ''; autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ]; }); torch = ignoreMissing prev.torch; triton = ignoreMissing prev.triton; } ); pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages { inherit python; }).overrideScope ( lib.composeManyExtensions [ inputs.pyproject-build-systems.overlays.default overlay exoOverlay buildSystemsOverlay linuxOverlay ] ); # mlx-cpu and mlx-cuda-13 both ship mlx/ site-packages files; keep first. # mlx-cpu/mlx-cuda-13 and nvidia-cudnn-cu12/cu13 ship overlapping files. venvCollisionPaths = lib.optionals pkgs.stdenv.hostPlatform.isLinux [ "lib/python3.13/site-packages/mlx*" "lib/python3.13/site-packages/nvidia*" ]; # Exclude bench deps from main env (bench has its own benchVenv) exoDeps = removeAttrs workspace.deps.default [ "exo-bench" ]; exoVenv = (pythonSet.mkVirtualEnv "exo-env" exoDeps).overrideAttrs { venvIgnoreCollisions = venvCollisionPaths; }; # Virtual environment with dev dependencies for testing testVenv = (pythonSet.mkVirtualEnv "exo-test-env" ( exoDeps // { exo = [ "dev" ]; # Include pytest, pytest-asyncio, pytest-env } )).overrideAttrs { venvIgnoreCollisions = venvCollisionPaths; }; mkPythonScript = name: path: pkgs.writeShellApplication { inherit name; runtimeInputs = [ exoVenv ]; runtimeEnv = { EXO_DASHBOARD_DIR = self'.packages.dashboard; EXO_RESOURCES_DIR = inputs.self + /resources; }; text = ''exec python ${path} "$@"''; }; benchVenv = pythonSet.mkVirtualEnv "exo-bench-env" { exo-bench = [ ]; }; mkBenchScript = name: path: pkgs.writeShellApplication { inherit name; runtimeInputs = [ benchVenv ]; text = ''exec python ${path} "$@"''; }; mkSimplePythonScript = name: path: pkgs.writeShellApplication { inherit name; runtimeInputs = [ pkgs.python313 ]; text = ''exec python ${path} "$@"''; }; exoPackage = pkgs.runCommand "exo" { nativeBuildInputs = [ pkgs.makeWrapper ]; } '' mkdir -p $out/bin # Create wrapper script makeWrapper ${exoVenv}/bin/exo $out/bin/exo \ --set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \ --set EXO_RESOURCES_DIR ${inputs.self + /resources} \ ${lib.optionalString pkgs.stdenv.hostPlatform.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"} ''; in { # Python package only available on macOS (requires MLX/Metal) packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin { exo = exoPackage; # Test environment for running pytest outside of Nix sandbox (needs GPU access) exo-test-env = testVenv; } // { exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py); exo-eval = mkBenchScript "exo-eval" (inputs.self + /bench/exo_eval.py); exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py); exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py); }; checks = { # Ruff linting (works on all platforms) lint = pkgs.runCommand "ruff-lint" { } '' export RUFF_CACHE_DIR="$TMPDIR/ruff-cache" ${pkgs.ruff}/bin/ruff check ${inputs.self} touch $out ''; # Hermetic basedpyright type checking typecheck = pkgs.runCommand "typecheck" { nativeBuildInputs = [ testVenv pkgs.basedpyright ]; } '' cd ${inputs.self} export HOME=$TMPDIR basedpyright --pythonpath ${testVenv}/bin/python touch $out ''; }; }; } ================================================ FILE: resources/image_model_cards/exolabs--FLUX.1-Kontext-dev-4bit.toml ================================================ model_id = "exolabs/FLUX.1-Kontext-dev-4bit" n_layers = 57 hidden_size = 1 supports_tensor = false tasks = ["ImageToImage"] family = "flux" quantization = "4bit" base_model = "FLUX.1 Kontext" capabilities = ["image_edit"] [storage_size] in_bytes = 15475325472 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 0 [[components]] component_name = "text_encoder_2" component_path = "text_encoder_2/" n_layers = 24 can_shard = false safetensors_index_filename = "model.safetensors.index.json" [components.storage_size] in_bytes = 9524621312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 57 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 5950704160 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--FLUX.1-Kontext-dev-8bit.toml ================================================ model_id = "exolabs/FLUX.1-Kontext-dev-8bit" n_layers = 57 hidden_size = 1 supports_tensor = false tasks = ["ImageToImage"] family = "flux" quantization = "8bit" base_model = "FLUX.1 Kontext" capabilities = ["image_edit"] [storage_size] in_bytes = 21426029632 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 0 [[components]] component_name = "text_encoder_2" component_path = "text_encoder_2/" n_layers = 24 can_shard = false safetensors_index_filename = "model.safetensors.index.json" [components.storage_size] in_bytes = 9524621312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 57 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 11901408320 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--FLUX.1-Kontext-dev.toml ================================================ model_id = "exolabs/FLUX.1-Kontext-dev" n_layers = 57 hidden_size = 1 supports_tensor = false tasks = ["ImageToImage"] family = "flux" quantization = "" base_model = "FLUX.1 Kontext" capabilities = ["image_edit"] [storage_size] in_bytes = 33327437952 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 0 [[components]] component_name = "text_encoder_2" component_path = "text_encoder_2/" n_layers = 24 can_shard = false safetensors_index_filename = "model.safetensors.index.json" [components.storage_size] in_bytes = 9524621312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 57 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 23802816640 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--FLUX.1-Krea-dev-4bit.toml ================================================ model_id = "exolabs/FLUX.1-Krea-dev-4bit" n_layers = 57 hidden_size = 1 supports_tensor = false tasks = ["TextToImage"] family = "flux" quantization = "4bit" base_model = "FLUX.1 Krea" capabilities = ["image_gen"] [storage_size] in_bytes = 15475325472 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 0 [[components]] component_name = "text_encoder_2" component_path = "text_encoder_2/" n_layers = 24 can_shard = false safetensors_index_filename = "model.safetensors.index.json" [components.storage_size] in_bytes = 9524621312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 57 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 5950704160 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--FLUX.1-Krea-dev-8bit.toml ================================================ model_id = "exolabs/FLUX.1-Krea-dev-8bit" n_layers = 57 hidden_size = 1 supports_tensor = false tasks = ["TextToImage"] family = "flux" quantization = "8bit" base_model = "FLUX.1 Krea" capabilities = ["image_gen"] [storage_size] in_bytes = 21426029632 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 0 [[components]] component_name = "text_encoder_2" component_path = "text_encoder_2/" n_layers = 24 can_shard = false safetensors_index_filename = "model.safetensors.index.json" [components.storage_size] in_bytes = 9524621312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 57 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 11901408320 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--FLUX.1-Krea-dev.toml ================================================ model_id = "exolabs/FLUX.1-Krea-dev" n_layers = 57 hidden_size = 1 supports_tensor = false tasks = ["TextToImage"] family = "flux" quantization = "" base_model = "FLUX.1 Krea" capabilities = ["image_gen"] [storage_size] in_bytes = 33327437952 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 0 [[components]] component_name = "text_encoder_2" component_path = "text_encoder_2/" n_layers = 24 can_shard = false safetensors_index_filename = "model.safetensors.index.json" [components.storage_size] in_bytes = 9524621312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 57 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 23802816640 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--FLUX.1-dev-4bit.toml ================================================ model_id = "exolabs/FLUX.1-dev-4bit" n_layers = 57 hidden_size = 1 supports_tensor = false tasks = ["TextToImage"] family = "flux" quantization = "4bit" base_model = "FLUX.1 Dev" capabilities = ["image_gen"] [storage_size] in_bytes = 15475325472 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 0 [[components]] component_name = "text_encoder_2" component_path = "text_encoder_2/" n_layers = 24 can_shard = false safetensors_index_filename = "model.safetensors.index.json" [components.storage_size] in_bytes = 9524621312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 57 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 5950704160 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--FLUX.1-dev-8bit.toml ================================================ model_id = "exolabs/FLUX.1-dev-8bit" n_layers = 57 hidden_size = 1 supports_tensor = false tasks = ["TextToImage"] family = "flux" quantization = "8bit" base_model = "FLUX.1 Dev" capabilities = ["image_gen"] [storage_size] in_bytes = 21426029632 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 0 [[components]] component_name = "text_encoder_2" component_path = "text_encoder_2/" n_layers = 24 can_shard = false safetensors_index_filename = "model.safetensors.index.json" [components.storage_size] in_bytes = 9524621312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 57 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 11901408320 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--FLUX.1-dev.toml ================================================ model_id = "exolabs/FLUX.1-dev" n_layers = 57 hidden_size = 1 supports_tensor = false tasks = ["TextToImage"] family = "flux" quantization = "" base_model = "FLUX.1 Dev" capabilities = ["image_gen"] [storage_size] in_bytes = 33327437952 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 0 [[components]] component_name = "text_encoder_2" component_path = "text_encoder_2/" n_layers = 24 can_shard = false safetensors_index_filename = "model.safetensors.index.json" [components.storage_size] in_bytes = 9524621312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 57 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 23802816640 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--FLUX.1-schnell-4bit.toml ================================================ model_id = "exolabs/FLUX.1-schnell-4bit" n_layers = 57 hidden_size = 1 supports_tensor = false tasks = ["TextToImage"] family = "flux" quantization = "4bit" base_model = "FLUX.1 Schnell" capabilities = ["image_gen"] [storage_size] in_bytes = 15470210592 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 0 [[components]] component_name = "text_encoder_2" component_path = "text_encoder_2/" n_layers = 24 can_shard = false safetensors_index_filename = "model.safetensors.index.json" [components.storage_size] in_bytes = 9524621312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 57 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 5945589280 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--FLUX.1-schnell-8bit.toml ================================================ model_id = "exolabs/FLUX.1-schnell-8bit" n_layers = 57 hidden_size = 1 supports_tensor = false tasks = ["TextToImage"] family = "flux" quantization = "8bit" base_model = "FLUX.1 Schnell" capabilities = ["image_gen"] [storage_size] in_bytes = 21415799872 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 0 [[components]] component_name = "text_encoder_2" component_path = "text_encoder_2/" n_layers = 24 can_shard = false safetensors_index_filename = "model.safetensors.index.json" [components.storage_size] in_bytes = 9524621312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 57 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 11891178560 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--FLUX.1-schnell.toml ================================================ model_id = "exolabs/FLUX.1-schnell" n_layers = 57 hidden_size = 1 supports_tensor = false tasks = ["TextToImage"] family = "flux" quantization = "" base_model = "FLUX.1 Schnell" capabilities = ["image_gen"] [storage_size] in_bytes = 33306978432 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 0 [[components]] component_name = "text_encoder_2" component_path = "text_encoder_2/" n_layers = 24 can_shard = false safetensors_index_filename = "model.safetensors.index.json" [components.storage_size] in_bytes = 9524621312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 57 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 23782357120 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--Qwen-Image-4bit.toml ================================================ model_id = "exolabs/Qwen-Image-4bit" n_layers = 60 hidden_size = 1 supports_tensor = false tasks = ["TextToImage"] uses_cfg = true family = "qwen-image" quantization = "4bit" base_model = "Qwen Image" capabilities = ["image_gen"] [storage_size] in_bytes = 26799533856 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 16584333312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 60 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 10215200544 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--Qwen-Image-8bit.toml ================================================ model_id = "exolabs/Qwen-Image-8bit" n_layers = 60 hidden_size = 1 supports_tensor = false tasks = ["TextToImage"] uses_cfg = true family = "qwen-image" quantization = "8bit" base_model = "Qwen Image" capabilities = ["image_gen"] [storage_size] in_bytes = 37014734400 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 16584333312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 60 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 20430401088 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--Qwen-Image-Edit-2509-4bit.toml ================================================ model_id = "exolabs/Qwen-Image-Edit-2509-4bit" n_layers = 60 hidden_size = 1 supports_tensor = false tasks = ["ImageToImage"] uses_cfg = true family = "qwen-image" quantization = "4bit" base_model = "Qwen Image Edit" capabilities = ["image_edit"] [storage_size] in_bytes = 26799533856 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 16584333312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 60 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 10215200544 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--Qwen-Image-Edit-2509-8bit.toml ================================================ model_id = "exolabs/Qwen-Image-Edit-2509-8bit" n_layers = 60 hidden_size = 1 supports_tensor = false tasks = ["ImageToImage"] uses_cfg = true family = "qwen-image" quantization = "8bit" base_model = "Qwen Image Edit" capabilities = ["image_edit"] [storage_size] in_bytes = 37014734400 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 16584333312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 60 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 20430401088 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--Qwen-Image-Edit-2509.toml ================================================ model_id = "exolabs/Qwen-Image-Edit-2509" n_layers = 60 hidden_size = 1 supports_tensor = false tasks = ["ImageToImage"] uses_cfg = true family = "qwen-image" quantization = "" base_model = "Qwen Image Edit" capabilities = ["image_edit"] [storage_size] in_bytes = 57445135488 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 16584333312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 60 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 40860802176 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/image_model_cards/exolabs--Qwen-Image.toml ================================================ model_id = "exolabs/Qwen-Image" n_layers = 60 hidden_size = 1 supports_tensor = false tasks = ["TextToImage"] uses_cfg = true family = "qwen-image" quantization = "" base_model = "Qwen Image" capabilities = ["image_gen"] [storage_size] in_bytes = 57445135488 [[components]] component_name = "text_encoder" component_path = "text_encoder/" n_layers = 12 can_shard = false [components.storage_size] in_bytes = 16584333312 [[components]] component_name = "transformer" component_path = "transformer/" n_layers = 60 can_shard = true safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json" [components.storage_size] in_bytes = 40860802176 [[components]] component_name = "vae" component_path = "vae/" can_shard = false [components.storage_size] in_bytes = 0 ================================================ FILE: resources/inference_model_cards/mlx-community--DeepSeek-V3.1-4bit.toml ================================================ model_id = "mlx-community/DeepSeek-V3.1-4bit" n_layers = 61 hidden_size = 7168 num_key_value_heads = 128 supports_tensor = true tasks = ["TextGeneration"] family = "deepseek" quantization = "4bit" base_model = "DeepSeek V3.1" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 405874409472 ================================================ FILE: resources/inference_model_cards/mlx-community--DeepSeek-V3.1-8bit.toml ================================================ model_id = "mlx-community/DeepSeek-V3.1-8bit" n_layers = 61 hidden_size = 7168 num_key_value_heads = 128 supports_tensor = true tasks = ["TextGeneration"] family = "deepseek" quantization = "8bit" base_model = "DeepSeek V3.1" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 765577920512 ================================================ FILE: resources/inference_model_cards/mlx-community--GLM-4.5-Air-8bit.toml ================================================ model_id = "mlx-community/GLM-4.5-Air-8bit" n_layers = 46 hidden_size = 4096 num_key_value_heads = 8 supports_tensor = false tasks = ["TextGeneration"] family = "glm" quantization = "8bit" base_model = "GLM 4.5 Air" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 122406567936 ================================================ FILE: resources/inference_model_cards/mlx-community--GLM-4.5-Air-bf16.toml ================================================ model_id = "mlx-community/GLM-4.5-Air-bf16" n_layers = 46 hidden_size = 4096 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "glm" quantization = "bf16" base_model = "GLM 4.5 Air" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 229780750336 ================================================ FILE: resources/inference_model_cards/mlx-community--GLM-4.7-4bit.toml ================================================ model_id = "mlx-community/GLM-4.7-4bit" n_layers = 91 hidden_size = 5120 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "glm" quantization = "4bit" base_model = "GLM 4.7" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 198556925568 ================================================ FILE: resources/inference_model_cards/mlx-community--GLM-4.7-6bit.toml ================================================ model_id = "mlx-community/GLM-4.7-6bit" n_layers = 91 hidden_size = 5120 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "glm" quantization = "6bit" base_model = "GLM 4.7" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 286737579648 ================================================ FILE: resources/inference_model_cards/mlx-community--GLM-4.7-8bit-gs32.toml ================================================ model_id = "mlx-community/GLM-4.7-8bit-gs32" n_layers = 91 hidden_size = 5120 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "glm" quantization = "8bit" base_model = "GLM 4.7" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 396963397248 ================================================ FILE: resources/inference_model_cards/mlx-community--GLM-4.7-Flash-4bit.toml ================================================ model_id = "mlx-community/GLM-4.7-Flash-4bit" n_layers = 47 hidden_size = 2048 num_key_value_heads = 20 supports_tensor = true tasks = ["TextGeneration"] family = "glm" quantization = "4bit" base_model = "GLM 4.7 Flash" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 19327352832 ================================================ FILE: resources/inference_model_cards/mlx-community--GLM-4.7-Flash-5bit.toml ================================================ model_id = "mlx-community/GLM-4.7-Flash-5bit" n_layers = 47 hidden_size = 2048 num_key_value_heads = 20 supports_tensor = true tasks = ["TextGeneration"] family = "glm" quantization = "5bit" base_model = "GLM 4.7 Flash" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 22548578304 ================================================ FILE: resources/inference_model_cards/mlx-community--GLM-4.7-Flash-6bit.toml ================================================ model_id = "mlx-community/GLM-4.7-Flash-6bit" n_layers = 47 hidden_size = 2048 num_key_value_heads = 20 supports_tensor = true tasks = ["TextGeneration"] family = "glm" quantization = "6bit" base_model = "GLM 4.7 Flash" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 26843545600 ================================================ FILE: resources/inference_model_cards/mlx-community--GLM-4.7-Flash-8bit.toml ================================================ model_id = "mlx-community/GLM-4.7-Flash-8bit" n_layers = 47 hidden_size = 2048 num_key_value_heads = 20 supports_tensor = true tasks = ["TextGeneration"] family = "glm" quantization = "8bit" base_model = "GLM 4.7 Flash" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 34359738368 ================================================ FILE: resources/inference_model_cards/mlx-community--GLM-5-8bit.toml ================================================ model_id = "mlx-community/GLM-5-8bit-MXFP8" n_layers = 78 hidden_size = 6144 num_key_value_heads = 64 supports_tensor = true tasks = ["TextGeneration"] family = "glm" quantization = "8bit" base_model = "GLM-5" capabilities = ["text", "thinking"] [storage_size] in_bytes = 790517400864 ================================================ FILE: resources/inference_model_cards/mlx-community--GLM-5-MXFP4-Q8.toml ================================================ model_id = "mlx-community/GLM-5-MXFP4-Q8" n_layers = 78 hidden_size = 6144 num_key_value_heads = 64 supports_tensor = true tasks = ["TextGeneration"] family = "glm" quantization = "MXFP4-Q8" base_model = "GLM-5" capabilities = ["text", "thinking"] [storage_size] in_bytes = 405478939008 ================================================ FILE: resources/inference_model_cards/mlx-community--GLM-5-bf16.toml ================================================ model_id = "mlx-community/GLM-5" n_layers = 78 hidden_size = 6144 num_key_value_heads = 64 supports_tensor = true tasks = ["TextGeneration"] family = "glm" quantization = "bf16" base_model = "GLM-5" capabilities = ["text", "thinking"] [storage_size] in_bytes = 1487822475264 ================================================ FILE: resources/inference_model_cards/mlx-community--Kimi-K2-Instruct-4bit.toml ================================================ model_id = "mlx-community/Kimi-K2-Instruct-4bit" n_layers = 61 hidden_size = 7168 num_key_value_heads = 64 supports_tensor = true tasks = ["TextGeneration"] family = "kimi" quantization = "4bit" base_model = "Kimi K2" capabilities = ["text"] [storage_size] in_bytes = 620622774272 ================================================ FILE: resources/inference_model_cards/mlx-community--Kimi-K2-Thinking.toml ================================================ model_id = "mlx-community/Kimi-K2-Thinking" n_layers = 61 hidden_size = 7168 num_key_value_heads = 64 supports_tensor = true tasks = ["TextGeneration"] family = "kimi" quantization = "" base_model = "Kimi K2" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 706522120192 ================================================ FILE: resources/inference_model_cards/mlx-community--Kimi-K2.5.toml ================================================ model_id = "mlx-community/Kimi-K2.5" n_layers = 61 hidden_size = 7168 num_key_value_heads = 64 supports_tensor = true tasks = ["TextGeneration"] family = "kimi" quantization = "" base_model = "Kimi K2.5" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 662498705408 ================================================ FILE: resources/inference_model_cards/mlx-community--Llama-3.1-Nemotron-70B-Instruct-HF-4bit.toml ================================================ model_id = "mlx-community/Llama-3.1-Nemotron-70B-Instruct-HF-4bit" n_layers = 80 hidden_size = 8192 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "4bit" base_model = "NVIDIA Llama-3.1-Nemotron-70B-Instruct" capabilities = ["text"] [storage_size] in_bytes = 39688355840 ================================================ FILE: resources/inference_model_cards/mlx-community--Llama-3.1-Nemotron-70B-Instruct-HF-8bit.toml ================================================ model_id = "mlx-community/Llama-3.1-Nemotron-70B-Instruct-HF-8bit" n_layers = 80 hidden_size = 8192 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "8bit" base_model = "NVIDIA Llama-3.1-Nemotron-70B-Instruct" capabilities = ["text"] [storage_size] in_bytes = 74964549632 ================================================ FILE: resources/inference_model_cards/mlx-community--Llama-3.1-Nemotron-70B-Instruct-HF-bf16.toml ================================================ model_id = "mlx-community/Llama-3.1-Nemotron-70B-Instruct-HF-bf16" n_layers = 80 hidden_size = 8192 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "bf16" base_model = "NVIDIA Llama-3.1-Nemotron-70B-Instruct" capabilities = ["text"] [storage_size] in_bytes = 141107412992 ================================================ FILE: resources/inference_model_cards/mlx-community--Llama-3.1-Nemotron-Nano-4B-v1.1-4bit.toml ================================================ model_id = "mlx-community/Llama-3.1-Nemotron-Nano-4B-v1.1-4bit" n_layers = 32 hidden_size = 3072 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "4bit" base_model = "NVIDIA Llama-3.1-Nemotron-Nano-4B-v1.1" capabilities = ["text"] [storage_size] in_bytes = 2538706944 ================================================ FILE: resources/inference_model_cards/mlx-community--Llama-3.1-Nemotron-Nano-4B-v1.1-8bit.toml ================================================ model_id = "mlx-community/Llama-3.1-Nemotron-Nano-4B-v1.1-8bit" n_layers = 32 hidden_size = 3072 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "8bit" base_model = "NVIDIA Llama-3.1-Nemotron-Nano-4B-v1.1" capabilities = ["text"] [storage_size] in_bytes = 4794980352 ================================================ FILE: resources/inference_model_cards/mlx-community--Llama-3.1-Nemotron-Nano-4B-v1.1-bf16.toml ================================================ model_id = "mlx-community/Llama-3.1-Nemotron-Nano-4B-v1.1-bf16" n_layers = 32 hidden_size = 3072 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "bf16" base_model = "NVIDIA Llama-3.1-Nemotron-Nano-4B-v1.1" capabilities = ["text"] [storage_size] in_bytes = 9025492992 ================================================ FILE: resources/inference_model_cards/mlx-community--Llama-3.2-1B-Instruct-4bit.toml ================================================ model_id = "mlx-community/Llama-3.2-1B-Instruct-4bit" n_layers = 16 hidden_size = 2048 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "4bit" base_model = "Llama 3.2 1B" capabilities = ["text"] [storage_size] in_bytes = 729808896 ================================================ FILE: resources/inference_model_cards/mlx-community--Llama-3.2-3B-Instruct-4bit.toml ================================================ model_id = "mlx-community/Llama-3.2-3B-Instruct-4bit" n_layers = 28 hidden_size = 3072 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "4bit" base_model = "Llama 3.2 3B" capabilities = ["text"] [storage_size] in_bytes = 1863319552 ================================================ FILE: resources/inference_model_cards/mlx-community--Llama-3.2-3B-Instruct-8bit.toml ================================================ model_id = "mlx-community/Llama-3.2-3B-Instruct-8bit" n_layers = 28 hidden_size = 3072 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "8bit" base_model = "Llama 3.2 3B" capabilities = ["text"] [storage_size] in_bytes = 3501195264 ================================================ FILE: resources/inference_model_cards/mlx-community--Llama-3.3-70B-Instruct-4bit.toml ================================================ model_id = "mlx-community/Llama-3.3-70B-Instruct-4bit" n_layers = 80 hidden_size = 8192 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "4bit" base_model = "Llama 3.3 70B" capabilities = ["text"] [storage_size] in_bytes = 40652242944 ================================================ FILE: resources/inference_model_cards/mlx-community--Llama-3.3-70B-Instruct-8bit.toml ================================================ model_id = "mlx-community/Llama-3.3-70B-Instruct-8bit" n_layers = 80 hidden_size = 8192 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "8bit" base_model = "Llama 3.3 70B" capabilities = ["text"] [storage_size] in_bytes = 76799803392 ================================================ FILE: resources/inference_model_cards/mlx-community--Meta-Llama-3.1-70B-Instruct-4bit.toml ================================================ model_id = "mlx-community/Meta-Llama-3.1-70B-Instruct-4bit" n_layers = 80 hidden_size = 8192 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "4bit" base_model = "Llama 3.1 70B" capabilities = ["text"] [storage_size] in_bytes = 40652242944 ================================================ FILE: resources/inference_model_cards/mlx-community--Meta-Llama-3.1-8B-Instruct-4bit.toml ================================================ model_id = "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit" n_layers = 32 hidden_size = 4096 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "4bit" base_model = "Llama 3.1 8B" capabilities = ["text"] [storage_size] in_bytes = 4637851648 ================================================ FILE: resources/inference_model_cards/mlx-community--Meta-Llama-3.1-8B-Instruct-8bit.toml ================================================ model_id = "mlx-community/Meta-Llama-3.1-8B-Instruct-8bit" n_layers = 32 hidden_size = 4096 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "8bit" base_model = "Llama 3.1 8B" capabilities = ["text"] [storage_size] in_bytes = 8954839040 ================================================ FILE: resources/inference_model_cards/mlx-community--Meta-Llama-3.1-8B-Instruct-bf16.toml ================================================ model_id = "mlx-community/Meta-Llama-3.1-8B-Instruct-bf16" n_layers = 32 hidden_size = 4096 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "bf16" base_model = "Llama 3.1 8B" capabilities = ["text"] [storage_size] in_bytes = 16882073600 ================================================ FILE: resources/inference_model_cards/mlx-community--MiniMax-M2.1-3bit.toml ================================================ model_id = "mlx-community/MiniMax-M2.1-3bit" n_layers = 61 hidden_size = 3072 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "minimax" quantization = "3bit" base_model = "MiniMax M2.1" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 100086644736 ================================================ FILE: resources/inference_model_cards/mlx-community--MiniMax-M2.1-8bit.toml ================================================ model_id = "mlx-community/MiniMax-M2.1-8bit" n_layers = 61 hidden_size = 3072 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "minimax" quantization = "8bit" base_model = "MiniMax M2.1" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 242986745856 ================================================ FILE: resources/inference_model_cards/mlx-community--MiniMax-M2.5-4bit.toml ================================================ model_id = "mlx-community/MiniMax-M2.5-4bit" n_layers = 62 hidden_size = 3072 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "minimax" quantization = "4bit" base_model = "MiniMax M2.5" capabilities = ["text", "thinking"] [storage_size] in_bytes = 128666664960 ================================================ FILE: resources/inference_model_cards/mlx-community--MiniMax-M2.5-6bit.toml ================================================ model_id = "mlx-community/MiniMax-M2.5-6bit" n_layers = 62 hidden_size = 3072 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "minimax" quantization = "6bit" base_model = "MiniMax M2.5" capabilities = ["text", "thinking"] [storage_size] in_bytes = 185826705408 ================================================ FILE: resources/inference_model_cards/mlx-community--MiniMax-M2.5-8bit.toml ================================================ model_id = "mlx-community/MiniMax-M2.5-8bit" n_layers = 62 hidden_size = 3072 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "minimax" quantization = "8bit" base_model = "MiniMax M2.5" capabilities = ["text", "thinking"] [storage_size] in_bytes = 242986745856 ================================================ FILE: resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-4Bit.toml ================================================ model_id = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-4Bit" n_layers = 52 hidden_size = 2688 supports_tensor = true tasks = ["TextGeneration"] family = "nemotron" quantization = "4bit" base_model = "NVIDIA Nemotron-3-Nano-30B-A3B" capabilities = ["text"] [storage_size] in_bytes = 17775342336 ================================================ FILE: resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-5Bit.toml ================================================ model_id = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-5Bit" n_layers = 52 hidden_size = 2688 supports_tensor = true tasks = ["TextGeneration"] family = "nemotron" quantization = "5bit" base_model = "NVIDIA Nemotron-3-Nano-30B-A3B" capabilities = ["text"] [storage_size] in_bytes = 21721476864 ================================================ FILE: resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-6Bit.toml ================================================ model_id = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-6Bit" n_layers = 52 hidden_size = 2688 supports_tensor = true tasks = ["TextGeneration"] family = "nemotron" quantization = "6bit" base_model = "NVIDIA Nemotron-3-Nano-30B-A3B" capabilities = ["text"] [storage_size] in_bytes = 25667611392 ================================================ FILE: resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-8Bit.toml ================================================ model_id = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-8Bit" n_layers = 52 hidden_size = 2688 supports_tensor = true tasks = ["TextGeneration"] family = "nemotron" quantization = "8bit" base_model = "NVIDIA Nemotron-3-Nano-30B-A3B" capabilities = ["text"] [storage_size] in_bytes = 33559880448 ================================================ FILE: resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-BF16.toml ================================================ model_id = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-BF16" n_layers = 52 hidden_size = 2688 supports_tensor = true tasks = ["TextGeneration"] family = "nemotron" quantization = "bf16" base_model = "NVIDIA Nemotron-3-Nano-30B-A3B" capabilities = ["text"] [storage_size] in_bytes = 63155889408 ================================================ FILE: resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-MXFP4.toml ================================================ model_id = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-MXFP4" n_layers = 52 hidden_size = 2688 supports_tensor = true tasks = ["TextGeneration"] family = "nemotron" quantization = "4bit" base_model = "NVIDIA Nemotron-3-Nano-30B-A3B" capabilities = ["text"] [storage_size] in_bytes = 16788808704 ================================================ FILE: resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4.toml ================================================ model_id = "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4" n_layers = 52 hidden_size = 2688 supports_tensor = true tasks = ["TextGeneration"] family = "nemotron" quantization = "4bit" base_model = "NVIDIA Nemotron-3-Nano-30B-A3B" capabilities = ["text"] [storage_size] in_bytes = 19323906944 ================================================ FILE: resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-Nano-9B-v2-4bits.toml ================================================ model_id = "mlx-community/NVIDIA-Nemotron-Nano-9B-v2-4bits" n_layers = 56 hidden_size = 4480 supports_tensor = true tasks = ["TextGeneration"] family = "nemotron" quantization = "4bit" base_model = "NVIDIA Nemotron-Nano-9B-v2" capabilities = ["text"] [storage_size] in_bytes = 5002791936 ================================================ FILE: resources/inference_model_cards/mlx-community--NVIDIA-Nemotron-Nano-9B-v2-6bit.toml ================================================ model_id = "mlx-community/NVIDIA-Nemotron-Nano-9B-v2-6bit" n_layers = 56 hidden_size = 4480 supports_tensor = true tasks = ["TextGeneration"] family = "nemotron" quantization = "6bit" base_model = "NVIDIA Nemotron-Nano-9B-v2" capabilities = ["text"] [storage_size] in_bytes = 7224298496 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-0.6B-4bit.toml ================================================ model_id = "mlx-community/Qwen3-0.6B-4bit" n_layers = 28 hidden_size = 1024 num_key_value_heads = 8 supports_tensor = false tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3 0.6B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 342884352 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-0.6B-8bit.toml ================================================ model_id = "mlx-community/Qwen3-0.6B-8bit" n_layers = 28 hidden_size = 1024 num_key_value_heads = 8 supports_tensor = false tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3 0.6B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 698351616 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-235B-A22B-Instruct-2507-4bit.toml ================================================ model_id = "mlx-community/Qwen3-235B-A22B-Instruct-2507-4bit" n_layers = 94 hidden_size = 4096 num_key_value_heads = 4 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3 235B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 141733920768 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-235B-A22B-Instruct-2507-8bit.toml ================================================ model_id = "mlx-community/Qwen3-235B-A22B-Instruct-2507-8bit" n_layers = 94 hidden_size = 4096 num_key_value_heads = 4 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3 235B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 268435456000 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-30B-A3B-4bit.toml ================================================ model_id = "mlx-community/Qwen3-30B-A3B-4bit" n_layers = 48 hidden_size = 2048 num_key_value_heads = 4 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3 30B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 17612931072 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-30B-A3B-8bit.toml ================================================ model_id = "mlx-community/Qwen3-30B-A3B-8bit" n_layers = 48 hidden_size = 2048 num_key_value_heads = 4 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3 30B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 33279705088 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-Coder-480B-A35B-Instruct-4bit.toml ================================================ model_id = "mlx-community/Qwen3-Coder-480B-A35B-Instruct-4bit" n_layers = 62 hidden_size = 6144 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3 Coder 480B" capabilities = ["text", "code"] [storage_size] in_bytes = 289910292480 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-Coder-480B-A35B-Instruct-8bit.toml ================================================ model_id = "mlx-community/Qwen3-Coder-480B-A35B-Instruct-8bit" n_layers = 62 hidden_size = 6144 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3 Coder 480B" capabilities = ["text", "code"] [storage_size] in_bytes = 579820584960 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-Coder-Next-4bit.toml ================================================ model_id = "mlx-community/Qwen3-Coder-Next-4bit" n_layers = 48 hidden_size = 2048 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3 Coder Next" capabilities = ["text", "code"] [storage_size] in_bytes = 45644286500 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-Coder-Next-5bit.toml ================================================ model_id = "mlx-community/Qwen3-Coder-Next-5bit" n_layers = 48 hidden_size = 2048 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "5bit" base_model = "Qwen3 Coder Next" capabilities = ["text", "code"] [storage_size] in_bytes = 57657697020 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-Coder-Next-6bit.toml ================================================ model_id = "mlx-community/Qwen3-Coder-Next-6bit" n_layers = 48 hidden_size = 2048 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "6bit" base_model = "Qwen3 Coder Next" capabilities = ["text", "code"] [storage_size] in_bytes = 68899327465 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-Coder-Next-8bit.toml ================================================ model_id = "mlx-community/Qwen3-Coder-Next-8bit" n_layers = 48 hidden_size = 2048 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3 Coder Next" capabilities = ["text", "code"] [storage_size] in_bytes = 89357758772 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-Coder-Next-bf16.toml ================================================ model_id = "mlx-community/Qwen3-Coder-Next-bf16" n_layers = 48 hidden_size = 2048 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "bf16" base_model = "Qwen3 Coder Next" capabilities = ["text", "code"] [storage_size] in_bytes = 157548627945 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Instruct-4bit.toml ================================================ model_id = "mlx-community/Qwen3-Next-80B-A3B-Instruct-4bit" n_layers = 48 hidden_size = 2048 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3 Next 80B" capabilities = ["text"] [storage_size] in_bytes = 46976204800 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Instruct-8bit.toml ================================================ model_id = "mlx-community/Qwen3-Next-80B-A3B-Instruct-8bit" n_layers = 48 hidden_size = 2048 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3 Next 80B" capabilities = ["text"] [storage_size] in_bytes = 88814387200 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-4bit.toml ================================================ model_id = "mlx-community/Qwen3-Next-80B-A3B-Thinking-4bit" n_layers = 48 hidden_size = 2048 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3 Next 80B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 47080074240 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-8bit.toml ================================================ model_id = "mlx-community/Qwen3-Next-80B-A3B-Thinking-8bit" n_layers = 48 hidden_size = 2048 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3 Next 80B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 88814387200 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-4bit.toml ================================================ model_id = "mlx-community/Qwen3.5-122B-A10B-4bit" n_layers = 48 hidden_size = 3072 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3.5 122B A10B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 69593314272 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-6bit.toml ================================================ model_id = "mlx-community/Qwen3.5-122B-A10B-6bit" n_layers = 48 hidden_size = 3072 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "6bit" base_model = "Qwen3.5 122B A10B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 100120675296 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-8bit.toml ================================================ model_id = "mlx-community/Qwen3.5-122B-A10B-8bit" n_layers = 48 hidden_size = 3072 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3.5 122B A10B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 130648036320 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3.5-122B-A10B-bf16.toml ================================================ model_id = "mlx-community/Qwen3.5-122B-A10B-bf16" n_layers = 48 hidden_size = 3072 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "bf16" base_model = "Qwen3.5 122B A10B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 245125640160 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3.5-27B-4bit.toml ================================================ model_id = "mlx-community/Qwen3.5-27B-4bit" n_layers = 64 hidden_size = 5120 num_key_value_heads = 4 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3.5 27B" capabilities = ["text", "thinking"] [storage_size] in_bytes = 16054266848 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3.5-27B-8bit.toml ================================================ model_id = "mlx-community/Qwen3.5-27B-8bit" n_layers = 64 hidden_size = 5120 num_key_value_heads = 4 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3.5 27B" capabilities = ["text", "thinking"] [storage_size] in_bytes = 29500943328 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3.5-2B-MLX-8bit.toml ================================================ model_id = "mlx-community/Qwen3.5-2B-MLX-8bit" n_layers = 24 hidden_size = 2048 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3.5 2B" capabilities = ["text", "thinking"] [storage_size] in_bytes = 2662787264 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3.5-35B-A3B-4bit.toml ================================================ model_id = "mlx-community/Qwen3.5-35B-A3B-4bit" n_layers = 40 hidden_size = 2048 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3.5 35B A3B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 20391405152 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3.5-35B-A3B-8bit.toml ================================================ model_id = "mlx-community/Qwen3.5-35B-A3B-8bit" n_layers = 40 hidden_size = 2048 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3.5 35B A3B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 37721130592 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-4bit.toml ================================================ model_id = "mlx-community/Qwen3.5-397B-A17B-4bit" n_layers = 60 hidden_size = 4096 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3.5 397B A17B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 223860768352 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-6bit.toml ================================================ model_id = "mlx-community/Qwen3.5-397B-A17B-6bit" n_layers = 60 hidden_size = 4096 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "6bit" base_model = "Qwen3.5 397B A17B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 322946674272 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3.5-397B-A17B-8bit.toml ================================================ model_id = "mlx-community/Qwen3.5-397B-A17B-8bit" n_layers = 60 hidden_size = 4096 num_key_value_heads = 2 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3.5 397B A17B" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 422032580192 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3.5-9B-4bit.toml ================================================ model_id = "mlx-community/Qwen3.5-9B-4bit" n_layers = 32 hidden_size = 4096 num_key_value_heads = 4 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3.5 9B" capabilities = ["text", "thinking"] [storage_size] in_bytes = 5950062560 ================================================ FILE: resources/inference_model_cards/mlx-community--Qwen3.5-9B-8bit.toml ================================================ model_id = "mlx-community/Qwen3.5-9B-8bit" n_layers = 32 hidden_size = 4096 num_key_value_heads = 4 supports_tensor = true tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3.5 9B" capabilities = ["text", "thinking"] [storage_size] in_bytes = 10426433504 ================================================ FILE: resources/inference_model_cards/mlx-community--Step-3.5-Flash-4bit.toml ================================================ model_id = "mlx-community/Step-3.5-Flash-4bit" n_layers = 45 hidden_size = 4096 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "step" quantization = "4bit" base_model = "Step 3.5 Flash" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 114572190076 ================================================ FILE: resources/inference_model_cards/mlx-community--Step-3.5-Flash-6bit.toml ================================================ model_id = "mlx-community/Step-3.5-Flash-6bit" n_layers = 45 hidden_size = 4096 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "step" quantization = "6bit" base_model = "Step 3.5 Flash" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 159039627774 ================================================ FILE: resources/inference_model_cards/mlx-community--Step-3.5-Flash-8Bit.toml ================================================ model_id = "mlx-community/Step-3.5-Flash-8Bit" n_layers = 45 hidden_size = 4096 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "step" quantization = "8bit" base_model = "Step 3.5 Flash" capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 209082699847 ================================================ FILE: resources/inference_model_cards/mlx-community--gpt-oss-120b-MXFP4-Q8.toml ================================================ model_id = "mlx-community/gpt-oss-120b-MXFP4-Q8" n_layers = 36 hidden_size = 2880 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "gpt-oss" quantization = "MXFP4-Q8" base_model = "GPT-OSS 120B" capabilities = ["text", "thinking"] [storage_size] in_bytes = 70652212224 ================================================ FILE: resources/inference_model_cards/mlx-community--gpt-oss-20b-MXFP4-Q8.toml ================================================ model_id = "mlx-community/gpt-oss-20b-MXFP4-Q8" n_layers = 24 hidden_size = 2880 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "gpt-oss" quantization = "MXFP4-Q8" base_model = "GPT-OSS 20B" capabilities = ["text", "thinking"] [storage_size] in_bytes = 12025908224 ================================================ FILE: resources/inference_model_cards/mlx-community--llama-3.3-70b-instruct-fp16.toml ================================================ model_id = "mlx-community/llama-3.3-70b-instruct-fp16" n_layers = 80 hidden_size = 8192 num_key_value_heads = 8 supports_tensor = true tasks = ["TextGeneration"] family = "llama" quantization = "fp16" base_model = "Llama 3.3 70B" capabilities = ["text"] [storage_size] in_bytes = 144383672320 ================================================ FILE: rust/exo_pyo3_bindings/Cargo.toml ================================================ [package] name = "exo_pyo3_bindings" version = { workspace = true } edition = { workspace = true } publish = false [lib] doctest = false path = "src/lib.rs" name = "exo_pyo3_bindings" # "cdylib" needed to produce shared library for Python to import # "rlib" needed for stub-gen to run crate-type = ["cdylib", "rlib"] [[bin]] path = "src/bin/stub_gen.rs" name = "stub_gen" doc = false [lints] workspace = true [dependencies] networking = { workspace = true } # interop pyo3 = { version = "0.27.2", features = [ # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13 # "nightly", # enables better-supported GIL integration "experimental-async", # async support in #[pyfunction] & #[pymethods] #"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation #"py-clone", # adding Clone-ing of `Py` without GIL (may cause panics - remove if panics happen) # "multiple-pymethods", # allows multiple #[pymethods] sections per class # integrations with other libraries # "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational", # "ordered-float", "rust_decimal", "smallvec", # "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time", "serde", ] } pyo3-stub-gen = { version = "0.17.2" } pyo3-async-runtimes = { version = "0.27.0", features = [ "attributes", "tokio-runtime", "testing", ] } pyo3-log = "0.13.2" # macro dependencies extend = { workspace = true } delegate = { workspace = true } # async runtime tokio = { workspace = true, features = ["full", "tracing"] } futures-lite = { workspace = true } # utility dependencies util = { workspace = true } # Tracing log = { workspace = true } env_logger = "0.11" # Networking libp2p = { workspace = true, features = ["full"] } pin-project = "1.1.10" ================================================ FILE: rust/exo_pyo3_bindings/README.md ================================================ TODO: do something here.... ================================================ FILE: rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi ================================================ # This file is automatically generated by pyo3_stub_gen # ruff: noqa: E501, F401 import builtins import typing @typing.final class AllQueuesFullError(builtins.Exception): def __new__(cls, *args: typing.Any) -> AllQueuesFullError: ... def __repr__(self) -> builtins.str: ... def __str__(self) -> builtins.str: ... @typing.final class Keypair: r""" Identity keypair of a node. """ @staticmethod def generate() -> Keypair: r""" Generate a new Ed25519 keypair. """ @staticmethod def from_bytes(bytes: bytes) -> Keypair: r""" Construct an Ed25519 keypair from secret key bytes """ def to_bytes(self) -> bytes: r""" Get the secret key bytes underlying the keypair """ def to_node_id(self) -> builtins.str: r""" Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`. """ @typing.final class MessageTooLargeError(builtins.Exception): def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ... def __repr__(self) -> builtins.str: ... def __str__(self) -> builtins.str: ... @typing.final class NetworkingHandle: def __new__(cls, identity: Keypair) -> NetworkingHandle: ... async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool: r""" Subscribe to a `GossipSub` topic. Returns `True` if the subscription worked. Returns `False` if we were already subscribed. """ async def gossipsub_unsubscribe(self, topic: builtins.str) -> builtins.bool: r""" Unsubscribes from a `GossipSub` topic. Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed. """ async def gossipsub_publish(self, topic: builtins.str, data: bytes) -> None: r""" Publishes a message with multiple topics to the `GossipSub` network. If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception. """ async def recv(self) -> PyFromSwarm: ... @typing.final class NoPeersSubscribedToTopicError(builtins.Exception): def __new__(cls, *args: typing.Any) -> NoPeersSubscribedToTopicError: ... def __repr__(self) -> builtins.str: ... def __str__(self) -> builtins.str: ... class PyFromSwarm: @typing.final class Connection(PyFromSwarm): __match_args__ = ("peer_id", "connected",) @property def peer_id(self) -> builtins.str: ... @property def connected(self) -> builtins.bool: ... def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> PyFromSwarm.Connection: ... @typing.final class Message(PyFromSwarm): __match_args__ = ("origin", "topic", "data",) @property def origin(self) -> builtins.str: ... @property def topic(self) -> builtins.str: ... @property def data(self) -> bytes: ... def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ... ... ================================================ FILE: rust/exo_pyo3_bindings/pyproject.toml ================================================ [build-system] requires = ["maturin>=1.0,<2.0"] build-backend = "maturin" [project] name = "exo_pyo3_bindings" version = "0.2.1" description = "Add your description here" readme = "README.md" authors = [ { name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" }, { name = "Evan Quiney", email = "evanev7@gmail.com" }, ] requires-python = ">=3.13" dependencies = [] [dependency-groups] dev = ["exo_pyo3_bindings", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"] [tool.maturin] #purelib = true #python-source = "python" module-name = "exo_pyo3_bindings" features = ["pyo3/extension-module", "pyo3/experimental-async"] [tool.pytest.ini_options] log_cli = true log_cli_level = "INFO" asyncio_mode = "auto" ================================================ FILE: rust/exo_pyo3_bindings/src/allow_threading.rs ================================================ //! SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await //! use pin_project::pin_project; use pyo3::prelude::*; use std::{ future::Future, pin::Pin, task::{Context, Poll}, }; /// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await #[pin_project] #[repr(transparent)] pub(crate) struct AllowThreads(#[pin] F); impl AllowThreads where Self: Future, { pub fn new(f: F) -> Self { Self(f) } } impl Future for AllowThreads where F: Future + Send, F::Output: Send, { type Output = F::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let waker = cx.waker(); Python::attach(|py| py.detach(|| self.project().0.poll(&mut Context::from_waker(waker)))) } } ================================================ FILE: rust/exo_pyo3_bindings/src/bin/stub_gen.rs ================================================ use pyo3_stub_gen::Result; fn main() -> Result<()> { env_logger::Builder::from_env(env_logger::Env::default().filter_or("RUST_LOG", "info")).init(); let stub = exo_pyo3_bindings::stub_info()?; stub.generate()?; Ok(()) } ================================================ FILE: rust/exo_pyo3_bindings/src/ident.rs ================================================ use crate::ext::ResultExt as _; use libp2p::identity::Keypair; use pyo3::types::{PyBytes, PyBytesMethods as _}; use pyo3::{Bound, PyResult, Python, pyclass, pymethods}; use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods}; /// Identity keypair of a node. #[gen_stub_pyclass] #[pyclass(name = "Keypair", frozen)] #[repr(transparent)] pub struct PyKeypair(pub Keypair); #[gen_stub_pymethods] #[pymethods] #[allow(clippy::needless_pass_by_value)] impl PyKeypair { /// Generate a new Ed25519 keypair. #[staticmethod] fn generate() -> Self { Self(Keypair::generate_ed25519()) } /// Construct an Ed25519 keypair from secret key bytes #[staticmethod] fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult { let mut bytes = Vec::from(bytes.as_bytes()); Ok(Self(Keypair::ed25519_from_bytes(&mut bytes).pyerr()?)) } /// Get the secret key bytes underlying the keypair fn to_bytes<'py>(&self, py: Python<'py>) -> PyResult> { let bytes = self .0 .clone() .try_into_ed25519() .pyerr()? .secret() .as_ref() .to_vec(); Ok(PyBytes::new(py, &bytes)) } /// Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`. fn to_node_id(&self) -> String { self.0.public().to_peer_id().to_base58() } } ================================================ FILE: rust/exo_pyo3_bindings/src/lib.rs ================================================ //! TODO: crate documentation //! //! this is here as a placeholder documentation //! //! mod allow_threading; mod ident; mod networking; use crate::ident::PyKeypair; use crate::networking::networking_submodule; use pyo3::prelude::PyModule; use pyo3::types::PyModuleMethods; use pyo3::{Bound, PyResult, pyclass, pymodule}; use pyo3_stub_gen::define_stub_info_gatherer; /// Namespace for all the constants used by this crate. pub(crate) mod r#const { pub const MPSC_CHANNEL_SIZE: usize = 1024; } /// Namespace for crate-wide extension traits/methods pub(crate) mod ext { use crate::allow_threading::AllowThreads; use extend::ext; use pyo3::exceptions::{PyConnectionError, PyRuntimeError}; use pyo3::types::PyBytes; use pyo3::{Py, PyErr, PyResult, Python}; use tokio::runtime::Runtime; use tokio::sync::mpsc; use tokio::sync::mpsc::error::TryRecvError; use tokio::task::JoinHandle; #[ext(pub, name = ByteArrayExt)] impl [u8] { fn pybytes(&self) -> Py { Python::attach(|py| PyBytes::new(py, self).unbind()) } } #[ext(pub, name = ResultExt)] impl Result where E: ToString, { fn pyerr(self) -> PyResult { self.map_err(|e| PyRuntimeError::new_err(e.to_string())) } } pub trait FutureExt: Future + Sized { /// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await fn allow_threads_py(self) -> AllowThreads where AllowThreads: Future, { AllowThreads::new(self) } } impl FutureExt for T {} #[ext(pub, name = PyErrExt)] impl PyErr { fn receiver_channel_closed() -> Self { PyConnectionError::new_err("Receiver channel closed unexpectedly") } } #[ext(pub, name = PyResultExt)] impl PyResult { fn write_unraisable(self) -> Option { Python::attach(|py| self.write_unraisable_with(py)) } fn write_unraisable_with(self, py: Python<'_>) -> Option { match self { Ok(v) => Some(v), Err(e) => { // write error back to python e.write_unraisable(py, None); None } } } } #[ext(pub, name = TokioRuntimeExt)] impl Runtime { fn spawn_with_scope(&self, py: Python<'_>, future: F) -> PyResult> where F: Future + Send + 'static, F::Output: Send + 'static, { let locals = pyo3_async_runtimes::tokio::get_current_locals(py)?; Ok(self.spawn(pyo3_async_runtimes::tokio::scope(locals, future))) } } #[ext(pub, name = TokioMpscSenderExt)] impl mpsc::Sender { /// Sends a value, waiting until there is capacity. /// /// A successful send occurs when it is determined that the other end of the /// channel has not hung up already. An unsuccessful send would be one where /// the corresponding receiver has already been closed. async fn send_py(&self, value: T) -> PyResult<()> { self.send(value) .await .map_err(|_| PyErr::receiver_channel_closed()) } } #[ext(pub, name = TokioMpscReceiverExt)] impl mpsc::Receiver { /// Receives the next value for this receiver. async fn recv_py(&mut self) -> PyResult { self.recv().await.ok_or_else(PyErr::receiver_channel_closed) } /// Receives at most `limit` values for this receiver and returns them. /// /// For `limit = 0`, an empty collection of messages will be returned immediately. /// For `limit > 0`, if there are no messages in the channel's queue this method /// will sleep until a message is sent. async fn recv_many_py(&mut self, limit: usize) -> PyResult> { // get updates from receiver channel let mut updates = Vec::with_capacity(limit); let received = self.recv_many(&mut updates, limit).await; // if we received zero items, then the channel was unexpectedly closed if limit != 0 && received == 0 { return Err(PyErr::receiver_channel_closed()); } Ok(updates) } /// Tries to receive the next value for this receiver. fn try_recv_py(&mut self) -> PyResult> { match self.try_recv() { Ok(v) => Ok(Some(v)), Err(TryRecvError::Empty) => Ok(None), Err(TryRecvError::Disconnected) => Err(PyErr::receiver_channel_closed()), } } } } /// A Python module implemented in Rust. The name of this function must match /// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to /// import the module. #[pymodule(name = "exo_pyo3_bindings")] fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> { // install logger pyo3_log::init(); let mut builder = tokio::runtime::Builder::new_multi_thread(); builder.enable_all(); pyo3_async_runtimes::tokio::init(builder); // TODO: for now this is all NOT a submodule, but figure out how to make the submodule system // work with maturin, where the types generate correctly, in the right folder, without // too many importing issues... m.add_class::()?; networking_submodule(m)?; // top-level constructs // TODO: ... Ok(()) } define_stub_info_gatherer!(stub_info); ================================================ FILE: rust/exo_pyo3_bindings/src/networking.rs ================================================ use std::pin::Pin; use std::sync::Arc; use crate::r#const::MPSC_CHANNEL_SIZE; use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _}; use crate::ext::{ResultExt as _, TokioMpscSenderExt as _}; use crate::ident::PyKeypair; use crate::networking::exception::{ PyAllQueuesFullError, PyMessageTooLargeError, PyNoPeersSubscribedToTopicError, }; use crate::pyclass; use futures_lite::{Stream, StreamExt as _}; use libp2p::gossipsub::PublishError; use networking::swarm::{FromSwarm, ToSwarm, create_swarm}; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::{PyModule, PyModuleMethods as _}; use pyo3::types::PyBytes; use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods}; use pyo3_stub_gen::derive::{ gen_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods, }; use tokio::sync::{Mutex, mpsc, oneshot}; mod exception { use pyo3::types::PyTuple; use pyo3::{exceptions::PyException, prelude::*}; use pyo3_stub_gen::derive::*; #[gen_stub_pyclass] #[pyclass(frozen, extends=PyException, name="NoPeersSubscribedToTopicError")] pub struct PyNoPeersSubscribedToTopicError {} impl PyNoPeersSubscribedToTopicError { const MSG: &'static str = "\ No peers are currently subscribed to receive messages on this topic. \ Wait for peers to subscribe or check your network connectivity."; /// Creates a new [ `PyErr` ] of this type. /// /// [`PyErr`] : https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html "PyErr in pyo3" pub(crate) fn new_err() -> PyErr { PyErr::new::(()) // TODO: check if this needs to be replaced??? } } #[gen_stub_pymethods] #[pymethods] impl PyNoPeersSubscribedToTopicError { #[new] #[pyo3(signature = (*args))] #[allow(unused_variables)] pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self { Self {} } fn __repr__(&self) -> String { format!("PeerId(\"{}\")", Self::MSG) } fn __str__(&self) -> String { Self::MSG.to_string() } } #[gen_stub_pyclass] #[pyclass(frozen, extends=PyException, name="AllQueuesFullError")] pub struct PyAllQueuesFullError {} impl PyAllQueuesFullError { const MSG: &'static str = "All libp2p peers are unresponsive, resend the message or reconnect."; /// Creates a new [ `PyErr` ] of this type. /// /// [`PyErr`] : https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html "PyErr in pyo3" pub(crate) fn new_err() -> PyErr { PyErr::new::(()) // TODO: check if this needs to be replaced??? } } #[gen_stub_pymethods] #[pymethods] impl PyAllQueuesFullError { #[new] #[pyo3(signature = (*args))] #[allow(unused_variables)] pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self { Self {} } fn __repr__(&self) -> String { format!("PeerId(\"{}\")", Self::MSG) } fn __str__(&self) -> String { Self::MSG.to_string() } } #[gen_stub_pyclass] #[pyclass(frozen, extends=PyException, name="MessageTooLargeError")] pub struct PyMessageTooLargeError {} impl PyMessageTooLargeError { const MSG: &'static str = "Gossipsub message exceeds max_transmit_size. Reduce prompt length or increase the limit."; pub(crate) fn new_err() -> PyErr { PyErr::new::(()) } } #[gen_stub_pymethods] #[pymethods] impl PyMessageTooLargeError { #[new] #[pyo3(signature = (*args))] #[allow(unused_variables)] pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self { Self {} } fn __repr__(&self) -> String { format!("MessageTooLargeError(\"{}\")", Self::MSG) } fn __str__(&self) -> String { Self::MSG.to_string() } } } #[gen_stub_pyclass] #[pyclass(name = "NetworkingHandle")] struct PyNetworkingHandle { // channels pub to_swarm: mpsc::Sender, pub swarm: Arc + Send>>>>, } #[gen_stub_pyclass_complex_enum] #[pyclass] enum PyFromSwarm { Connection { peer_id: String, connected: bool, }, Message { origin: String, topic: String, data: Py, }, } impl From for PyFromSwarm { fn from(value: FromSwarm) -> Self { match value { FromSwarm::Discovered { peer_id } => Self::Connection { peer_id: peer_id.to_base58(), connected: true, }, FromSwarm::Expired { peer_id } => Self::Connection { peer_id: peer_id.to_base58(), connected: false, }, FromSwarm::Message { from, topic, data } => Self::Message { origin: from.to_base58(), topic: topic, data: data.pybytes(), }, } } } #[gen_stub_pymethods] #[pymethods] impl PyNetworkingHandle { // NOTE: `async fn`s here that use `.await` will wrap the future in `.allow_threads_py()` // immediately beforehand to release the interpreter. // SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await // ---- Lifecycle management methods ---- #[new] fn py_new(identity: Bound<'_, PyKeypair>) -> PyResult { // create communication channels let (to_swarm, from_client) = mpsc::channel(MPSC_CHANNEL_SIZE); // get identity let identity = identity.borrow().0.clone(); // create networking swarm (within tokio context!! or it crashes) let _guard = pyo3_async_runtimes::tokio::get_runtime().enter(); let swarm = create_swarm(identity, from_client).pyerr()?.into_stream(); Ok(Self { swarm: Arc::new(Mutex::new(swarm)), to_swarm, }) } #[gen_stub(skip)] fn recv<'py>(&'py self, py: Python<'py>) -> PyResult> { let swarm = Arc::clone(&self.swarm); pyo3_async_runtimes::tokio::future_into_py(py, async move { swarm .try_lock() .map_err(|_| PyRuntimeError::new_err("called recv twice concurrently"))? .next() .await .ok_or(PyErr::receiver_channel_closed()) .map(PyFromSwarm::from) }) } // ---- Gossipsub management methods ---- /// Subscribe to a `GossipSub` topic. /// /// Returns `True` if the subscription worked. Returns `False` if we were already subscribed. async fn gossipsub_subscribe(&self, topic: String) -> PyResult { let (tx, rx) = oneshot::channel(); // send off request to subscribe self.to_swarm .send_py(ToSwarm::Subscribe { topic, result_sender: tx, }) .allow_threads_py() // allow-threads-aware async call .await?; // wait for response & return any errors rx.allow_threads_py() // allow-threads-aware async call .await .map_err(|_| PyErr::receiver_channel_closed())? .pyerr() } /// Unsubscribes from a `GossipSub` topic. /// /// Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed. async fn gossipsub_unsubscribe(&self, topic: String) -> PyResult { let (tx, rx) = oneshot::channel(); // send off request to unsubscribe self.to_swarm .send_py(ToSwarm::Unsubscribe { topic, result_sender: tx, }) .allow_threads_py() // allow-threads-aware async call .await?; // wait for response & convert any errors rx.allow_threads_py() // allow-threads-aware async call .await .map_err(|_| PyErr::receiver_channel_closed()) } /// Publishes a message with multiple topics to the `GossipSub` network. /// /// If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception. async fn gossipsub_publish(&self, topic: String, data: Py) -> PyResult<()> { let (tx, rx) = oneshot::channel(); // send off request to subscribe let data = Python::attach(|py| Vec::from(data.as_bytes(py))); self.to_swarm .send_py(ToSwarm::Publish { topic, data, result_sender: tx, }) .allow_threads_py() // allow-threads-aware async call .await?; // wait for response & return any errors => ignore messageID for now!!! let _ = rx .allow_threads_py() // allow-threads-aware async call .await .map_err(|_| PyErr::receiver_channel_closed())? .map_err(|e| match e { PublishError::AllQueuesFull(_) => PyAllQueuesFullError::new_err(), PublishError::MessageTooLarge => PyMessageTooLargeError::new_err(), PublishError::NoPeersSubscribedToTopic => { PyNoPeersSubscribedToTopicError::new_err() } e => PyRuntimeError::new_err(e.to_string()), })?; Ok(()) } } pyo3_stub_gen::inventory::submit! { gen_methods_from_python! { r#" class PyNetworkingHandle: async def recv() -> PyFromSwarm: ... "# } } pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; Ok(()) } ================================================ FILE: rust/exo_pyo3_bindings/tests/dummy.rs ================================================ #[cfg(test)] mod tests { use core::mem::drop; use core::option::Option::Some; use core::time::Duration; use tokio; use tokio::sync::mpsc; #[tokio::test] async fn test_drop_channel() { struct Ping; let (tx, mut rx) = mpsc::channel::(10); let _ = tokio::spawn(async move { println!("TASK: entered"); loop { tokio::select! { result = rx.recv() => { match result { Some(_) => { println!("TASK: pinged"); } None => { println!("TASK: closing channel"); break; } } } _ = tokio::time::sleep(Duration::from_secs_f32(0.1)) => { println!("TASK: heartbeat"); } } } println!("TASK: exited"); }); let tx2 = tx.clone(); tokio::time::sleep(Duration::from_secs_f32(0.11)).await; tx.send(Ping).await.expect("Should not fail"); drop(tx); tokio::time::sleep(Duration::from_secs_f32(0.11)).await; tx2.send(Ping).await.expect("Should not fail"); drop(tx2); tokio::time::sleep(Duration::from_secs_f32(0.11)).await; } } ================================================ FILE: rust/exo_pyo3_bindings/tests/test_python.py ================================================ import asyncio import pytest from exo_pyo3_bindings import ( Keypair, NetworkingHandle, NoPeersSubscribedToTopicError, PyFromSwarm, ) @pytest.mark.asyncio async def test_sleep_on_multiple_items() -> None: print("PYTHON: starting handle") h = NetworkingHandle(Keypair.generate()) rt = asyncio.create_task(_await_recv(h)) # sleep for 4 ticks for i in range(4): await asyncio.sleep(1) try: await h.gossipsub_publish("topic", b"somehting or other") except NoPeersSubscribedToTopicError as e: print("caught it", e) async def _await_recv(h: NetworkingHandle): while True: event = await h.recv() match event: case PyFromSwarm.Connection() as c: print(f"PYTHON: connection update: {c}") case PyFromSwarm.Message() as m: print(f"PYTHON: message: {m}") ================================================ FILE: rust/networking/Cargo.toml ================================================ [package] name = "networking" version = { workspace = true } edition = { workspace = true } publish = false [lib] doctest = false name = "networking" path = "src/lib.rs" [lints] workspace = true [dependencies] # datastructures either = { workspace = true } # macro dependencies extend = { workspace = true } delegate = { workspace = true } # async async-stream = { workspace = true } futures-lite = { workspace = true } futures-timer = { workspace = true } tokio = { workspace = true, features = ["full"] } # utility dependencies util = { workspace = true } tracing-subscriber = { version = "0.3.19", features = [ "default", "env-filter", ] } keccak-const = { workspace = true } # tracing/logging log = { workspace = true } # networking libp2p = { workspace = true, features = ["full"] } pin-project = "1.1.10" ================================================ FILE: rust/networking/examples/chatroom.rs ================================================ use futures_lite::StreamExt; use libp2p::identity; use networking::swarm; use networking::swarm::{FromSwarm, ToSwarm}; use tokio::sync::{mpsc, oneshot}; use tokio::{io, io::AsyncBufReadExt as _}; use tracing_subscriber::EnvFilter; use tracing_subscriber::filter::LevelFilter; #[tokio::main] async fn main() { let _ = tracing_subscriber::fmt() .with_env_filter(EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into())) .try_init(); let (to_swarm, from_client) = mpsc::channel(20); // Configure swarm let mut swarm = swarm::create_swarm(identity::Keypair::generate_ed25519(), from_client) .expect("Swarm creation failed") .into_stream(); // Create a Gossipsub topic & subscribe let (tx, rx) = oneshot::channel(); _ = to_swarm .send(ToSwarm::Subscribe { topic: "test-net".to_string(), result_sender: tx, }) .await .expect("should send"); // Read full lines from stdin let mut stdin = io::BufReader::new(io::stdin()).lines(); println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub"); tokio::task::spawn(async move { rx.await .expect("tx not dropped") .expect("subscribe shouldn't fail"); loop { if let Ok(Some(line)) = stdin.next_line().await { let (tx, rx) = oneshot::channel(); if let Err(e) = to_swarm .send(swarm::ToSwarm::Publish { topic: "test-net".to_string(), data: line.as_bytes().to_vec(), result_sender: tx, }) .await { println!("Send error: {e:?}"); return; }; match rx.await { Ok(Err(e)) => println!("Publish error: {e:?}"), Err(e) => println!("Publish error: {e:?}"), Ok(_) => {} } } } }); // Kick it off loop { // on gossipsub outgoing match swarm.next().await { // on gossipsub incoming Some(FromSwarm::Discovered { peer_id }) => { println!("\n\nconnected to {peer_id}\n\n") } Some(FromSwarm::Expired { peer_id }) => { println!("\n\ndisconnected from {peer_id}\n\n") } Some(FromSwarm::Message { from, topic, data }) => { println!("{topic}/{from}:\n{}", String::from_utf8_lossy(&data)) } None => {} } } } ================================================ FILE: rust/networking/src/RESEARCH_NOTES.txt ================================================ https://github.com/ml-explore/mlx/commit/3fe98bacc7640d857acf3539f1d21b47a32e5609 ^raw sockets distributed -> `` -> https://newosxbook.com/code/xnu-3247.1.106/bsd/net/ndrv.h.auto.html --> header file for a networking component found in the macOS kernel (XNU) that defines structures for network device driver registration, specifically the ndrv_demux_desc and ndrv_protocol_desc structures used for demultiplexing protocol data at the network interface level. It specifies how to describe protocol data, such as an Ethernet type or a SNAP header, and how to associate these descriptions with a specific protocol family to receive matching packets. --> Used to bind an NDRV socket so that packets that match given protocol demux descriptions can be received. --> An NDRV socket is a special kind of socket in the Darwin/macOS operating system's XNU kernel, used for low-level network packet manipulation and binding to specific protocols for packet processing. It allows user-space applications or drivers to directly write Layer 2 (L2) network packets or interact with the network stack at a lower level, often by binding to protocol descriptors like the ndrv_protocol_desc. This type of socket is used for functions such as capturing and injecting packets, especially in network infrastructure software like routers or for kernel-level network monitoring and security tools. --> also called PF_NDRV sockets --> https://newosxbook.com/bonus/vol1ch16.html ----> they are conceptually similar to https://scapy.disruptivelabs.in/networking/socket-interface PF_RAW or PF_PACKET https://stackoverflow.com/questions/17169298/af-packet-on-osx ^AF_PACKET duplicates the packets as soon as it receives them from the physical layer (for incoming packets) or just before sending them out to the physical layer (for outgoing packets). -> this is on Linux only ^it doesn't exist on OS X so you can use /dev/bpfX (Berkeley Packet Filter) for sniffing https://www.unix.com/man_page/mojave/4/ip/ ^OS X manpages for IP https://developer.apple.com/documentation/kernel/implementing_drivers_system_extensions_and_kexts ^driver kit, system extensions & kexts for macOS ---- To set up a Linux system to use a Thunderbolt connection as a network device, connect the two computers with a Thunderbolt cable, load the thunderbolt-net kernel module (usually automatic but modprobe is an option for manual loading), and then the operating system will create virtual Ethernet interfaces (e.g., thunderbolt0) for networking. You can then use standard tools like ifconfig or your desktop environment's network manager to configure these new interfaces for a link-local network. --> https://gist.github.com/geosp/80fbd39e617b7d1d9421683df4ea224a ----> here is a guide on how to set up thunderbolt-ethernet on linux ----> I may be able to steal the thunderbolt-net code ideas to implement a kernel module for MacOS https://chatgpt.com/s/t_68af8e41a8548191993281a014f846a7 ^GPT discussion about making socket interface https://chatgpt.com/s/t_68afb798a85c8191973c02a0fa7a48a3 --> link-local address,,?? https://chatgpt.com/s/t_68afb02987e08191b2b0044d3667ece2 ^GPT discussion about accessing TB on MacOS low level interactions -------------------------------- https://www.intel.com/content/www/us/en/support/articles/000098893/software.html ^Thunderbolt Share & Thunderbolt Networking Mode => intel's equivalent of thunderbolt bridge --------------------------------- https://www.zerotier.com/blog/how-zerotier-eliminated-kernel-extensions-on-macos/ -->fake ethernet devices on MacOS -> omg??? we can detect thunderbolt bridge, then bind to it, then re-expose it as fake ethernet?? -->ps: https://chatgpt.com/s/t_68afb2b25fb881919526763fb5d7359c, AF/PF_NDRV are one and the same!!! -->https://github.com/zerotier/ZeroTierOne/blob/dev/osdep/MacEthernetTapAgent.c ================================================ FILE: rust/networking/src/discovery.rs ================================================ use crate::ext::MultiaddrExt; use delegate::delegate; use either::Either; use futures_lite::FutureExt; use futures_timer::Delay; use libp2p::core::transport::PortUse; use libp2p::core::{ConnectedPoint, Endpoint}; use libp2p::swarm::behaviour::ConnectionEstablished; use libp2p::swarm::dial_opts::DialOpts; use libp2p::swarm::{ CloseConnection, ConnectionClosed, ConnectionDenied, ConnectionHandler, ConnectionHandlerSelect, ConnectionId, FromSwarm, NetworkBehaviour, THandler, THandlerInEvent, THandlerOutEvent, ToSwarm, dummy, }; use libp2p::{Multiaddr, PeerId, identity, mdns}; use std::collections::{BTreeSet, HashMap}; use std::convert::Infallible; use std::io; use std::net::IpAddr; use std::task::{Context, Poll}; use std::time::Duration; use util::wakerdeque::WakerDeque; const RETRY_CONNECT_INTERVAL: Duration = Duration::from_secs(5); mod managed { use libp2p::swarm::NetworkBehaviour; use libp2p::{identity, mdns, ping}; use std::io; use std::time::Duration; const MDNS_RECORD_TTL: Duration = Duration::from_secs(2_500); const MDNS_QUERY_INTERVAL: Duration = Duration::from_secs(1_500); const PING_TIMEOUT: Duration = Duration::from_millis(2_500); const PING_INTERVAL: Duration = Duration::from_millis(2_500); #[derive(NetworkBehaviour)] pub struct Behaviour { mdns: mdns::tokio::Behaviour, ping: ping::Behaviour, } impl Behaviour { pub fn new(keypair: &identity::Keypair) -> io::Result { Ok(Self { mdns: mdns_behaviour(keypair)?, ping: ping_behaviour(), }) } } fn mdns_behaviour(keypair: &identity::Keypair) -> io::Result { use mdns::{Config, tokio}; // mDNS config => enable IPv6 let mdns_config = Config { ttl: MDNS_RECORD_TTL, query_interval: MDNS_QUERY_INTERVAL, // enable_ipv6: true, // TODO: for some reason, TCP+mDNS don't work well with ipv6?? figure out how to make work ..Default::default() }; let mdns_behaviour = tokio::Behaviour::new(mdns_config, keypair.public().to_peer_id()); Ok(mdns_behaviour?) } fn ping_behaviour() -> ping::Behaviour { ping::Behaviour::new( ping::Config::new() .with_timeout(PING_TIMEOUT) .with_interval(PING_INTERVAL), ) } } /// Events for when a listening connection is truly established and truly closed. #[derive(Debug, Clone)] pub enum Event { ConnectionEstablished { peer_id: PeerId, connection_id: ConnectionId, remote_ip: IpAddr, remote_tcp_port: u16, }, ConnectionClosed { peer_id: PeerId, connection_id: ConnectionId, remote_ip: IpAddr, remote_tcp_port: u16, }, } /// Discovery behavior that wraps mDNS to produce truly discovered durable peer-connections. /// /// The behaviour operates as such: /// 1) All true (listening) connections/disconnections are tracked, emitting corresponding events /// to the swarm. /// 1) mDNS discovered/expired peers are tracked; discovered but not connected peers are dialed /// immediately, and expired but connected peers are disconnected from immediately. /// 2) Every fixed interval: discovered but not connected peers are dialed, and expired but /// connected peers are disconnected from. pub struct Behaviour { // state-tracking for managed behaviors & mDNS-discovered peers managed: managed::Behaviour, mdns_discovered: HashMap>, retry_delay: Delay, // retry interval // pending events to emmit => waker-backed Deque to control polling pending_events: WakerDeque>, } impl Behaviour { pub fn new(keypair: &identity::Keypair) -> io::Result { Ok(Self { managed: managed::Behaviour::new(keypair)?, mdns_discovered: HashMap::new(), retry_delay: Delay::new(RETRY_CONNECT_INTERVAL), pending_events: WakerDeque::new(), }) } fn dial(&mut self, peer_id: PeerId, addr: Multiaddr) { self.pending_events.push_back(ToSwarm::Dial { opts: DialOpts::peer_id(peer_id).addresses(vec![addr]).build(), }) } fn close_connection(&mut self, peer_id: PeerId, connection: ConnectionId) { // push front to make this IMMEDIATE self.pending_events.push_front(ToSwarm::CloseConnection { peer_id, connection: CloseConnection::One(connection), }) } fn handle_mdns_discovered(&mut self, peers: Vec<(PeerId, Multiaddr)>) { for (p, ma) in peers { self.dial(p, ma.clone()); // always connect // get peer's multi-addresses or insert if missing let Some(mas) = self.mdns_discovered.get_mut(&p) else { self.mdns_discovered.insert(p, BTreeSet::from([ma])); continue; }; // multiaddress should never already be present - else something has gone wrong let is_new_addr = mas.insert(ma); assert!(is_new_addr, "cannot discover a discovered peer"); } } fn handle_mdns_expired(&mut self, peers: Vec<(PeerId, Multiaddr)>) { for (p, ma) in peers { // at this point, we *must* have the peer let mas = self .mdns_discovered .get_mut(&p) .expect("nonexistent peer cannot expire"); // at this point, we *must* have the multiaddress let was_present = mas.remove(&ma); assert!(was_present, "nonexistent multiaddress cannot expire"); // if empty, remove the peer-id entirely if mas.is_empty() { self.mdns_discovered.remove(&p); } } } fn on_connection_established( &mut self, peer_id: PeerId, connection_id: ConnectionId, remote_ip: IpAddr, remote_tcp_port: u16, ) { // send out connected event self.pending_events .push_back(ToSwarm::GenerateEvent(Event::ConnectionEstablished { peer_id, connection_id, remote_ip, remote_tcp_port, })); } fn on_connection_closed( &mut self, peer_id: PeerId, connection_id: ConnectionId, remote_ip: IpAddr, remote_tcp_port: u16, ) { // send out disconnected event self.pending_events .push_back(ToSwarm::GenerateEvent(Event::ConnectionClosed { peer_id, connection_id, remote_ip, remote_tcp_port, })); } } impl NetworkBehaviour for Behaviour { type ConnectionHandler = ConnectionHandlerSelect>; type ToSwarm = Event; // simply delegate to underlying mDNS behaviour delegate! { to self.managed { fn handle_pending_inbound_connection(&mut self, connection_id: ConnectionId, local_addr: &Multiaddr, remote_addr: &Multiaddr) -> Result<(), ConnectionDenied>; fn handle_pending_outbound_connection(&mut self, connection_id: ConnectionId, maybe_peer: Option, addresses: &[Multiaddr], effective_role: Endpoint) -> Result, ConnectionDenied>; } } fn handle_established_inbound_connection( &mut self, connection_id: ConnectionId, peer: PeerId, local_addr: &Multiaddr, remote_addr: &Multiaddr, ) -> Result, ConnectionDenied> { Ok(ConnectionHandler::select( dummy::ConnectionHandler, self.managed.handle_established_inbound_connection( connection_id, peer, local_addr, remote_addr, )?, )) } #[allow(clippy::needless_question_mark)] fn handle_established_outbound_connection( &mut self, connection_id: ConnectionId, peer: PeerId, addr: &Multiaddr, role_override: Endpoint, port_use: PortUse, ) -> Result, ConnectionDenied> { Ok(ConnectionHandler::select( dummy::ConnectionHandler, self.managed.handle_established_outbound_connection( connection_id, peer, addr, role_override, port_use, )?, )) } fn on_connection_handler_event( &mut self, peer_id: PeerId, connection_id: ConnectionId, event: THandlerOutEvent, ) { match event { Either::Left(ev) => libp2p::core::util::unreachable(ev), Either::Right(ev) => { self.managed .on_connection_handler_event(peer_id, connection_id, ev) } } } // hook into these methods to drive behavior fn on_swarm_event(&mut self, event: FromSwarm) { self.managed.on_swarm_event(event); // let mDNS handle swarm events // handle swarm events to update internal state: match event { FromSwarm::ConnectionEstablished(ConnectionEstablished { peer_id, connection_id, endpoint, .. }) => { let remote_address = match endpoint { ConnectedPoint::Dialer { address, .. } => address, ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr, }; if let Some((ip, port)) = remote_address.try_to_tcp_addr() { // handle connection established event which is filtered correctly self.on_connection_established(peer_id, connection_id, ip, port) } } FromSwarm::ConnectionClosed(ConnectionClosed { peer_id, connection_id, endpoint, .. }) => { let remote_address = match endpoint { ConnectedPoint::Dialer { address, .. } => address, ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr, }; if let Some((ip, port)) = remote_address.try_to_tcp_addr() { // handle connection closed event which is filtered correctly self.on_connection_closed(peer_id, connection_id, ip, port) } } // since we are running TCP/IP transport layer, we are assuming that // no address changes can occur, hence encountering one is a fatal error FromSwarm::AddressChange(a) => { unreachable!("unhandlable: address change encountered: {:?}", a) } _ => {} } } fn poll(&mut self, cx: &mut Context) -> Poll>> { // delegate to managed behaviors for any behaviors they need to perform match self.managed.poll(cx) { Poll::Ready(ToSwarm::GenerateEvent(e)) => { match e { // handle discovered and expired events from mDNS managed::BehaviourEvent::Mdns(e) => match e.clone() { mdns::Event::Discovered(peers) => { self.handle_mdns_discovered(peers); } mdns::Event::Expired(peers) => { self.handle_mdns_expired(peers); } }, // handle ping events => if error then disconnect managed::BehaviourEvent::Ping(e) => { if let Err(_) = e.result { self.close_connection(e.peer, e.connection.clone()) } } } // since we just consumed an event, we should immediately wake just in case // there are more events to come where that came from cx.waker().wake_by_ref(); } // forward any other mDNS event to the swarm or its connection handler(s) Poll::Ready(e) => { return Poll::Ready( e.map_out(|_| unreachable!("events returning to swarm already handled")) .map_in(Either::Right), ); } Poll::Pending => {} } // retry connecting to all mDNS peers periodically (fails safely if already connected) if self.retry_delay.poll(cx).is_ready() { for (p, mas) in self.mdns_discovered.clone() { for ma in mas { self.dial(p, ma) } } self.retry_delay.reset(RETRY_CONNECT_INTERVAL) // reset timeout } // send out any pending events from our own service if let Some(e) = self.pending_events.pop_front(cx) { return Poll::Ready(e.map_in(Either::Left)); } // wait for pending events Poll::Pending } } ================================================ FILE: rust/networking/src/lib.rs ================================================ //! TODO: crate documentation //! //! this is here as a placeholder documentation //! //! pub mod discovery; pub mod swarm; /// Namespace for all the type/trait aliases used by this crate. pub(crate) mod alias { use std::error::Error; pub type AnyError = Box; pub type AnyResult = Result; } /// Namespace for crate-wide extension traits/methods pub(crate) mod ext { use extend::ext; use libp2p::Multiaddr; use libp2p::multiaddr::Protocol; use std::net::IpAddr; #[ext(pub, name = MultiaddrExt)] impl Multiaddr { /// If the multiaddress corresponds to a TCP address, extracts it fn try_to_tcp_addr(&self) -> Option<(IpAddr, u16)> { let mut ps = self.into_iter(); let ip = if let Some(p) = ps.next() { match p { Protocol::Ip4(ip) => IpAddr::V4(ip), Protocol::Ip6(ip) => IpAddr::V6(ip), _ => return None, } } else { return None; }; let Some(Protocol::Tcp(port)) = ps.next() else { return None; }; Some((ip, port)) } } } ================================================ FILE: rust/networking/src/swarm.rs ================================================ use std::pin::Pin; use crate::swarm::transport::tcp_transport; use crate::{alias, discovery}; pub use behaviour::{Behaviour, BehaviourEvent}; use futures_lite::{Stream, StreamExt}; use libp2p::{PeerId, SwarmBuilder, gossipsub, identity, swarm::SwarmEvent}; use tokio::sync::{mpsc, oneshot}; /// The current version of the network: this prevents devices running different versions of the /// software from interacting with each other. /// /// TODO: right now this is a hardcoded constant; figure out what the versioning semantics should /// even be, and how to inject the right version into this config/initialization. E.g. should /// this be passed in as a parameter? What about rapidly changing versions in debug builds? /// this is all VERY very hard to figure out and needs to be mulled over as a team. pub const NETWORK_VERSION: &[u8] = b"v0.0.1"; pub const OVERRIDE_VERSION_ENV_VAR: &str = "EXO_LIBP2P_NAMESPACE"; // Uses oneshot senders to emulate function calling apis while avoiding requiring unique ownership // of the Swarm. pub enum ToSwarm { Unsubscribe { topic: String, result_sender: oneshot::Sender, }, Subscribe { topic: String, result_sender: oneshot::Sender>, }, Publish { topic: String, data: Vec, result_sender: oneshot::Sender>, }, } pub enum FromSwarm { Message { from: PeerId, topic: String, data: Vec, }, Discovered { peer_id: PeerId, }, Expired { peer_id: PeerId, }, } pub struct Swarm { swarm: libp2p::Swarm, from_client: mpsc::Receiver, } impl Swarm { pub fn into_stream(self) -> Pin + Send>> { let Swarm { mut swarm, mut from_client, } = self; let stream = async_stream::stream! { loop { tokio::select! { msg = from_client.recv() => { let Some(msg) = msg else { break }; on_message(&mut swarm, msg); } event = swarm.next() => { let Some(event) = event else { break }; if let Some(item) = filter_swarm_event(event) { yield item; } } } } }; Box::pin(stream) } } fn on_message(swarm: &mut libp2p::Swarm, message: ToSwarm) { match message { ToSwarm::Subscribe { topic, result_sender, } => { let result = swarm .behaviour_mut() .gossipsub .subscribe(&gossipsub::IdentTopic::new(topic)); _ = result_sender.send(result); } ToSwarm::Unsubscribe { topic, result_sender, } => { let result = swarm .behaviour_mut() .gossipsub .unsubscribe(&gossipsub::IdentTopic::new(topic)); _ = result_sender.send(result); } ToSwarm::Publish { topic, data, result_sender, } => { let result = swarm .behaviour_mut() .gossipsub .publish(gossipsub::IdentTopic::new(topic), data); _ = result_sender.send(result); } } } fn filter_swarm_event(event: SwarmEvent) -> Option { match event { SwarmEvent::Behaviour(BehaviourEvent::Gossipsub(gossipsub::Event::Message { message: gossipsub::Message { source: Some(peer_id), topic, data, .. }, .. })) => Some(FromSwarm::Message { from: peer_id, topic: topic.into_string(), data, }), SwarmEvent::Behaviour(BehaviourEvent::Discovery( discovery::Event::ConnectionEstablished { peer_id, .. }, )) => Some(FromSwarm::Discovered { peer_id }), SwarmEvent::Behaviour(BehaviourEvent::Discovery(discovery::Event::ConnectionClosed { peer_id, .. })) => Some(FromSwarm::Expired { peer_id }), _ => None, } } /// Create and configure a swarm which listens to all ports on OS pub fn create_swarm( keypair: identity::Keypair, from_client: mpsc::Receiver, ) -> alias::AnyResult { let mut swarm = SwarmBuilder::with_existing_identity(keypair) .with_tokio() .with_other_transport(tcp_transport)? .with_behaviour(Behaviour::new)? .build(); // Listen on all interfaces and whatever port the OS assigns swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?; Ok(Swarm { swarm, from_client }) } mod transport { use crate::alias; use crate::swarm::{NETWORK_VERSION, OVERRIDE_VERSION_ENV_VAR}; use futures_lite::{AsyncRead, AsyncWrite}; use keccak_const::Sha3_256; use libp2p::core::muxing; use libp2p::core::transport::Boxed; use libp2p::pnet::{PnetError, PnetOutput}; use libp2p::{PeerId, Transport, identity, noise, pnet, yamux}; use std::{env, sync::LazyLock}; /// Key used for networking's private network; parametrized on the [`NETWORK_VERSION`]. /// See [`pnet_upgrade`] for more. static PNET_PRESHARED_KEY: LazyLock<[u8; 32]> = LazyLock::new(|| { let builder = Sha3_256::new().update(b"exo_discovery_network"); if let Ok(var) = env::var(OVERRIDE_VERSION_ENV_VAR) { let bytes = var.into_bytes(); builder.update(&bytes) } else { builder.update(NETWORK_VERSION) } .finalize() }); /// Make the Swarm run on a private network, as to not clash with public libp2p nodes and /// also different-versioned instances of this same network. /// This is implemented as an additional "upgrade" ontop of existing [`libp2p::Transport`] layers. async fn pnet_upgrade( socket: TSocket, _: impl Sized, ) -> Result, PnetError> where TSocket: AsyncRead + AsyncWrite + Send + Unpin + 'static, { use pnet::{PnetConfig, PreSharedKey}; PnetConfig::new(PreSharedKey::new(*PNET_PRESHARED_KEY)) .handshake(socket) .await } /// TCP/IP transport layer configuration. pub fn tcp_transport( keypair: &identity::Keypair, ) -> alias::AnyResult> { use libp2p::{ core::upgrade::Version, tcp::{Config, tokio}, }; // `TCP_NODELAY` enabled => avoid latency let tcp_config = Config::default().nodelay(true); // V1 + lazy flushing => 0-RTT negotiation let upgrade_version = Version::V1Lazy; // Noise is faster than TLS + we don't care much for security let noise_config = noise::Config::new(keypair)?; // Use default Yamux config for multiplexing let yamux_config = yamux::Config::default(); // Create new Tokio-driven TCP/IP transport layer let base_transport = tokio::Transport::new(tcp_config) .and_then(pnet_upgrade) .upgrade(upgrade_version) .authenticate(noise_config) .multiplex(yamux_config); // Return boxed transport (to flatten complex type) Ok(base_transport.boxed()) } } mod behaviour { use crate::{alias, discovery}; use libp2p::swarm::NetworkBehaviour; use libp2p::{gossipsub, identity}; /// Behavior of the Swarm which composes all desired behaviors: /// Right now its just [`discovery::Behaviour`] and [`gossipsub::Behaviour`]. #[derive(NetworkBehaviour)] pub struct Behaviour { pub discovery: discovery::Behaviour, pub gossipsub: gossipsub::Behaviour, } impl Behaviour { pub fn new(keypair: &identity::Keypair) -> alias::AnyResult { Ok(Self { discovery: discovery::Behaviour::new(keypair)?, gossipsub: gossipsub_behaviour(keypair), }) } } fn gossipsub_behaviour(keypair: &identity::Keypair) -> gossipsub::Behaviour { use gossipsub::{ConfigBuilder, MessageAuthenticity, ValidationMode}; // build a gossipsub network behaviour // => signed message authenticity + strict validation mode means the message-ID is // automatically provided by gossipsub w/out needing to provide custom message-ID function gossipsub::Behaviour::new( MessageAuthenticity::Signed(keypair.clone()), ConfigBuilder::default() .max_transmit_size(8 * 1024 * 1024) .validation_mode(ValidationMode::Strict) .build() .expect("the configuration should always be valid"), ) .expect("creating gossipsub behavior should always work") } } ================================================ FILE: rust/networking/tests/dummy.rs ================================================ // maybe this will hold test in the future...?? #[cfg(test)] mod tests { #[test] fn does_nothing() {} } ================================================ FILE: rust/parts.nix ================================================ { inputs, ... }: { perSystem = { inputs', pkgs, lib, ... }: let # Fenix nightly toolchain with all components rustToolchain = inputs'.fenix.packages.stable.withComponents [ "cargo" "rustc" "clippy" "rustfmt" "rust-src" "rust-analyzer" ]; # Crane with fenix toolchain craneLib = (inputs.crane.mkLib pkgs).overrideToolchain rustToolchain; # Source filtering - only include rust/ directory and root Cargo files # This ensures changes to Python/docs/etc don't trigger Rust rebuilds src = lib.cleanSourceWith { src = inputs.self; filter = path: type: let baseName = builtins.baseNameOf path; parentDir = builtins.dirOf path; inRustDir = (lib.hasInfix "/rust/" path) || (lib.hasSuffix "/rust" parentDir) || (baseName == "rust" && type == "directory"); isRootCargoFile = (baseName == "Cargo.toml" || baseName == "Cargo.lock") && (builtins.dirOf path == toString inputs.self); in isRootCargoFile || (inRustDir && (craneLib.filterCargoSources path type || lib.hasSuffix ".toml" path || lib.hasSuffix ".md" path)); }; # Common arguments for all Rust builds commonArgs = { inherit src; pname = "exo-rust"; version = "0.0.1"; strictDeps = true; nativeBuildInputs = [ pkgs.pkg-config pkgs.python313 # Required for pyo3-build-config ]; buildInputs = [ pkgs.openssl pkgs.python313 # Required for pyo3 tests ]; OPENSSL_NO_VENDOR = "1"; # Required for pyo3 tests to find libpython LD_LIBRARY_PATH = lib.makeLibraryPath [ pkgs.python313 ]; }; # Build dependencies once for caching cargoArtifacts = craneLib.buildDepsOnly ( commonArgs // { cargoExtraArgs = "--workspace"; } ); in { # Export toolchain for use in treefmt and devShell options.rust = { toolchain = lib.mkOption { type = lib.types.package; default = rustToolchain; description = "The Rust toolchain to use"; }; }; config = { packages = { # Python bindings wheel via maturin exo_pyo3_bindings = craneLib.buildPackage ( commonArgs // { inherit cargoArtifacts; pname = "exo_pyo3_bindings"; nativeBuildInputs = commonArgs.nativeBuildInputs ++ [ pkgs.maturin ]; buildPhaseCargoCommand = '' maturin build \ --release \ --manylinux off \ --manifest-path rust/exo_pyo3_bindings/Cargo.toml \ --features "pyo3/extension-module,pyo3/experimental-async" \ --interpreter ${pkgs.python313}/bin/python \ --out dist ''; # Don't use crane's default install behavior doNotPostBuildInstallCargoBinaries = true; installPhaseCommand = '' mkdir -p $out cp dist/*.whl $out/ ''; } ); }; checks = { # Full workspace build (all crates) cargo-build = craneLib.buildPackage ( commonArgs // { inherit cargoArtifacts; cargoExtraArgs = "--workspace"; } ); # Run tests with nextest cargo-nextest = craneLib.cargoNextest ( commonArgs // { inherit cargoArtifacts; cargoExtraArgs = "--workspace"; } ); # Build documentation cargo-doc = craneLib.cargoDoc ( commonArgs // { inherit cargoArtifacts; cargoExtraArgs = "--workspace"; } ); }; }; }; } ================================================ FILE: rust/util/Cargo.toml ================================================ [package] name = "util" version = { workspace = true } edition = { workspace = true } publish = false [lib] doctest = false name = "util" path = "src/lib.rs" [lints] workspace = true [dependencies] ================================================ FILE: rust/util/src/lib.rs ================================================ pub mod wakerdeque; ================================================ FILE: rust/util/src/wakerdeque.rs ================================================ use std::collections::VecDeque; use std::fmt::{Debug, Formatter}; use std::task::{Context, Waker}; /// A wrapper around [`VecDeque`] which wakes (if it can) on any `push_*` methods, /// and updates the internally stored waker by consuming [`Context`] on any `pop_*` methods. pub struct WakerDeque { waker: Option, deque: VecDeque, } impl Debug for WakerDeque { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { self.deque.fmt(f) } } impl WakerDeque { pub fn new() -> Self { Self { waker: None, deque: VecDeque::new(), } } fn update(&mut self, cx: &mut Context<'_>) { self.waker = Some(cx.waker().clone()); } fn wake(&mut self) { let Some(ref mut w) = self.waker else { return }; w.wake_by_ref(); self.waker = None; } pub fn pop_front(&mut self, cx: &mut Context<'_>) -> Option { self.update(cx); self.deque.pop_front() } pub fn pop_back(&mut self, cx: &mut Context<'_>) -> Option { self.update(cx); self.deque.pop_back() } pub fn push_front(&mut self, value: T) { self.wake(); self.deque.push_front(value); } pub fn push_back(&mut self, value: T) { self.wake(); self.deque.push_back(value); } } ================================================ FILE: scripts/fetch_kv_heads.py ================================================ #!/usr/bin/env python3 """Fetch num_key_value_heads from HuggingFace config.json and update TOML model cards. Usage: # Update only cards missing num_key_value_heads uv run python scripts/fetch_kv_heads.py --missing # Update all cards (overwrite existing values) uv run python scripts/fetch_kv_heads.py --all """ from __future__ import annotations import argparse import json import sys import urllib.request from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path import tomlkit CARDS_DIR = ( Path(__file__).resolve().parent.parent / "resources" / "inference_model_cards" ) MAX_WORKERS = 5 def fetch_kv_heads(model_id: str) -> int | None: """Fetch num_key_value_heads from HuggingFace config.json.""" url = f"https://huggingface.co/{model_id}/raw/main/config.json" try: with urllib.request.urlopen(url, timeout=15) as resp: config = json.loads(resp.read()) except Exception as e: print(f" ERROR fetching {url}: {e}", file=sys.stderr) return None for source in [config, config.get("text_config", {})]: if "num_key_value_heads" in source: return int(source["num_key_value_heads"]) return None def update_toml(path: Path, kv_heads: int) -> bool: """Insert or update num_key_value_heads in a TOML file. Returns True if changed.""" content = path.read_text() doc = tomlkit.parse(content) if doc.get("num_key_value_heads") == kv_heads: return False # Insert after hidden_size if adding for the first time if "num_key_value_heads" not in doc: new_doc = tomlkit.document() for key, value in doc.items(): new_doc[key] = value if key == "hidden_size": new_doc["num_key_value_heads"] = kv_heads path.write_text(tomlkit.dumps(new_doc)) else: doc["num_key_value_heads"] = kv_heads path.write_text(tomlkit.dumps(doc)) return True def process_card(path: Path) -> tuple[str, str]: """Fetch and update a single card. Returns (filename, status).""" content = path.read_text() doc = tomlkit.parse(content) model_id = doc.get("model_id") if not model_id: return path.name, "SKIP (no model_id)" kv_heads = fetch_kv_heads(str(model_id)) if kv_heads is None: return path.name, "FAILED" changed = update_toml(path, kv_heads) return path.name, f"{kv_heads} ({'UPDATED' if changed else 'UNCHANGED'})" def main(): parser = argparse.ArgumentParser( description="Fetch num_key_value_heads from HuggingFace and update TOML cards." ) group = parser.add_mutually_exclusive_group(required=True) group.add_argument( "--all", action="store_true", help="Update all model cards (overwrite existing values)", ) group.add_argument( "--missing", action="store_true", help="Only update cards missing num_key_value_heads", ) args = parser.parse_args() toml_files = sorted(CARDS_DIR.glob("*.toml")) if not toml_files: print(f"No TOML files found in {CARDS_DIR}", file=sys.stderr) sys.exit(1) to_process = [] skipped = 0 for path in toml_files: if args.missing and "num_key_value_heads" in path.read_text(): skipped += 1 continue to_process.append(path) updated = 0 failed = 0 with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool: futures = {pool.submit(process_card, path): path for path in to_process} for future in as_completed(futures): name, status = future.result() print(f" {name}: {status}") if "UPDATED" in status: updated += 1 elif "FAILED" in status: failed += 1 print(f"\nDone: {updated} updated, {skipped} skipped, {failed} failed") if __name__ == "__main__": main() ================================================ FILE: src/exo/__init__.py ================================================ ================================================ FILE: src/exo/__main__.py ================================================ from __future__ import annotations import sys from collections.abc import Sequence from multiprocessing import freeze_support from typing import Final from exo.main import main INLINE_CODE_FLAG: Final[str] = "-c" def _maybe_run_inline_code(argv: Sequence[str]) -> bool: """ Reproduce the bare minimum of Python's `-c` flag so multiprocessing helper processes (for example the resource tracker) can execute. """ try: flag_index = argv.index(INLINE_CODE_FLAG) except ValueError: return False code_index = flag_index + 1 if code_index >= len(argv): return False inline_code = argv[code_index] sys.argv = ["-c", *argv[code_index + 1 :]] namespace: dict[str, object] = {"__name__": "__main__"} exec(inline_code, namespace, namespace) return True if __name__ == "__main__": if _maybe_run_inline_code(sys.argv): sys.exit(0) freeze_support() main() ================================================ FILE: src/exo/api/__init__.py ================================================ ================================================ FILE: src/exo/api/adapters/__init__.py ================================================ """API adapters for different API formats (Claude, OpenAI Responses, etc.).""" ================================================ FILE: src/exo/api/adapters/chat_completions.py ================================================ """OpenAI Chat Completions API adapter for converting requests/responses.""" import time from collections.abc import AsyncGenerator from typing import Any from exo.api.types import ( ChatCompletionChoice, ChatCompletionMessage, ChatCompletionMessageText, ChatCompletionRequest, ChatCompletionResponse, ErrorInfo, ErrorResponse, FinishReason, Logprobs, LogprobsContentItem, StreamingChoiceResponse, ToolCall, Usage, ) from exo.shared.types.chunks import ( ErrorChunk, PrefillProgressChunk, TokenChunk, ToolCallChunk, ) from exo.shared.types.common import CommandId from exo.shared.types.text_generation import ( InputMessage, TextGenerationTaskParams, resolve_reasoning_params, ) def chat_request_to_text_generation( request: ChatCompletionRequest, ) -> TextGenerationTaskParams: instructions: str | None = None input_messages: list[InputMessage] = [] chat_template_messages: list[dict[str, Any]] = [] for msg in request.messages: # Normalize content to string content: str if msg.content is None: content = "" elif isinstance(msg.content, str): content = msg.content elif isinstance(msg.content, ChatCompletionMessageText): content = msg.content.text else: # List of ChatCompletionMessageText content = "\n".join(item.text for item in msg.content) # Extract system message as instructions if msg.role == "system": if instructions is None: instructions = content else: # Append additional system messages instructions = f"{instructions}\n{content}" chat_template_messages.append({"role": "system", "content": content}) else: # Skip messages with no meaningful content if ( msg.content is None and msg.reasoning_content is None and msg.tool_calls is None ): continue if msg.role in ("user", "assistant", "developer"): input_messages.append(InputMessage(role=msg.role, content=content)) # Build full message dict for chat template (preserves tool_calls etc.) # Normalize content for model_dump msg_copy = msg.model_copy(update={"content": content}) dumped: dict[str, Any] = msg_copy.model_dump(exclude_none=True) chat_template_messages.append(dumped) resolved_effort, resolved_thinking = resolve_reasoning_params( request.reasoning_effort, request.enable_thinking ) return TextGenerationTaskParams( model=request.model, input=input_messages if input_messages else [InputMessage(role="user", content="")], instructions=instructions, max_output_tokens=request.max_tokens, temperature=request.temperature, top_p=request.top_p, top_k=request.top_k, stop=request.stop, seed=request.seed, stream=request.stream, tools=request.tools, reasoning_effort=resolved_effort, enable_thinking=resolved_thinking, chat_template_messages=chat_template_messages if chat_template_messages else None, logprobs=request.logprobs or False, top_logprobs=request.top_logprobs, min_p=request.min_p, repetition_penalty=request.repetition_penalty, repetition_context_size=request.repetition_context_size, ) def chunk_to_response( chunk: TokenChunk, command_id: CommandId ) -> ChatCompletionResponse: """Convert a TokenChunk to a streaming ChatCompletionResponse.""" # Build logprobs if available logprobs: Logprobs | None = None if chunk.logprob is not None: logprobs = Logprobs( content=[ LogprobsContentItem( token=chunk.text, logprob=chunk.logprob, top_logprobs=chunk.top_logprobs or [], ) ] ) if chunk.is_thinking: delta = ChatCompletionMessage(role="assistant", reasoning_content=chunk.text) else: delta = ChatCompletionMessage(role="assistant", content=chunk.text) return ChatCompletionResponse( id=command_id, created=int(time.time()), model=chunk.model, choices=[ StreamingChoiceResponse( index=0, delta=delta, logprobs=logprobs, finish_reason=chunk.finish_reason, ) ], ) async def generate_chat_stream( command_id: CommandId, chunk_stream: AsyncGenerator[ PrefillProgressChunk | ErrorChunk | ToolCallChunk | TokenChunk, None ], ) -> AsyncGenerator[str, None]: """Generate Chat Completions API streaming events from chunks.""" last_usage: Usage | None = None async for chunk in chunk_stream: match chunk: case PrefillProgressChunk(): # Use SSE comment so third-party clients ignore it yield f": prefill_progress {chunk.model_dump_json()}\n\n" case ErrorChunk(): error_response = ErrorResponse( error=ErrorInfo( message=chunk.error_message or "Internal server error", type="InternalServerError", code=500, ) ) yield f"data: {error_response.model_dump_json()}\n\n" yield "data: [DONE]\n\n" return case ToolCallChunk(): last_usage = chunk.usage or last_usage tool_call_deltas = [ ToolCall( id=tool.id, index=i, function=tool, ) for i, tool in enumerate(chunk.tool_calls) ] tool_response = ChatCompletionResponse( id=command_id, created=int(time.time()), model=chunk.model, choices=[ StreamingChoiceResponse( index=0, delta=ChatCompletionMessage( role="assistant", tool_calls=tool_call_deltas, ), finish_reason="tool_calls", ) ], usage=last_usage, ) yield f"data: {tool_response.model_dump_json()}\n\n" yield "data: [DONE]\n\n" return case TokenChunk(): last_usage = chunk.usage or last_usage chunk_response = chunk_to_response(chunk, command_id) if chunk.finish_reason is not None: chunk_response = chunk_response.model_copy( update={"usage": last_usage} ) yield f"data: {chunk_response.model_dump_json()}\n\n" if chunk.finish_reason is not None: yield "data: [DONE]\n\n" async def collect_chat_response( command_id: CommandId, chunk_stream: AsyncGenerator[ ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None ], ) -> AsyncGenerator[str]: # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason """Collect all token chunks and return a single ChatCompletionResponse.""" text_parts: list[str] = [] thinking_parts: list[str] = [] tool_calls: list[ToolCall] = [] logprobs_content: list[LogprobsContentItem] = [] model: str | None = None finish_reason: FinishReason | None = None error_message: str | None = None last_usage: Usage | None = None async for chunk in chunk_stream: match chunk: case PrefillProgressChunk(): continue case ErrorChunk(): error_message = chunk.error_message or "Internal server error" break case TokenChunk(): if model is None: model = chunk.model last_usage = chunk.usage or last_usage if chunk.is_thinking: thinking_parts.append(chunk.text) else: text_parts.append(chunk.text) if chunk.logprob is not None: logprobs_content.append( LogprobsContentItem( token=chunk.text, logprob=chunk.logprob, top_logprobs=chunk.top_logprobs or [], ) ) if chunk.finish_reason is not None: finish_reason = chunk.finish_reason case ToolCallChunk(): if model is None: model = chunk.model last_usage = chunk.usage or last_usage tool_calls.extend( ToolCall( id=tool.id, index=i, function=tool, ) for i, tool in enumerate(chunk.tool_calls) ) finish_reason = chunk.finish_reason if error_message is not None: raise ValueError(error_message) combined_text = "".join(text_parts) combined_thinking = "".join(thinking_parts) if thinking_parts else None assert model is not None yield ChatCompletionResponse( id=command_id, created=int(time.time()), model=model, choices=[ ChatCompletionChoice( index=0, message=ChatCompletionMessage( role="assistant", content=combined_text, reasoning_content=combined_thinking, tool_calls=tool_calls if tool_calls else None, ), logprobs=Logprobs(content=logprobs_content) if logprobs_content else None, finish_reason=finish_reason, ) ], usage=last_usage, ).model_dump_json() return ================================================ FILE: src/exo/api/adapters/claude.py ================================================ """Claude Messages API adapter for converting requests/responses.""" import json import re from collections.abc import AsyncGenerator from typing import Any from exo.api.types import FinishReason, Usage from exo.api.types.claude_api import ( ClaudeContentBlock, ClaudeContentBlockDeltaEvent, ClaudeContentBlockStartEvent, ClaudeContentBlockStopEvent, ClaudeInputJsonDelta, ClaudeMessageDelta, ClaudeMessageDeltaEvent, ClaudeMessageDeltaUsage, ClaudeMessagesRequest, ClaudeMessagesResponse, ClaudeMessageStart, ClaudeMessageStartEvent, ClaudeMessageStopEvent, ClaudeStopReason, ClaudeTextBlock, ClaudeTextDelta, ClaudeThinkingBlock, ClaudeThinkingDelta, ClaudeToolResultBlock, ClaudeToolUseBlock, ClaudeUsage, ) from exo.shared.types.chunks import ( ErrorChunk, PrefillProgressChunk, TokenChunk, ToolCallChunk, ) from exo.shared.types.common import CommandId from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams def finish_reason_to_claude_stop_reason( finish_reason: FinishReason | None, ) -> ClaudeStopReason | None: """Map OpenAI finish_reason to Claude stop_reason.""" if finish_reason is None: return None mapping: dict[FinishReason, ClaudeStopReason] = { "stop": "end_turn", "length": "max_tokens", "tool_calls": "tool_use", "content_filter": "end_turn", "function_call": "tool_use", } return mapping.get(finish_reason, "end_turn") def _extract_tool_result_text(block: ClaudeToolResultBlock) -> str: """Extract plain text from a tool_result content field.""" if block.content is None: return "" if isinstance(block.content, str): return block.content return "".join(sub_block.text for sub_block in block.content) # Matches "x-anthropic-billing-header: ...;" (with optional trailing newline) # or similar telemetry headers that change every request and break KV prefix caching. _VOLATILE_HEADER_RE = re.compile(r"^x-anthropic-[^\n]*;\n?", re.MULTILINE) def _strip_volatile_headers(text: str) -> str: """Remove Anthropic billing/telemetry headers from system prompt text. Claude Code prepends headers like 'x-anthropic-billing-header: cc_version=...; cc_entrypoint=...; cch=...;' that contain per-request content hashes. These change every request and break KV prefix caching (the prefix diverges at ~20 tokens instead of matching thousands of conversation tokens). """ return _VOLATILE_HEADER_RE.sub("", text) def claude_request_to_text_generation( request: ClaudeMessagesRequest, ) -> TextGenerationTaskParams: # Handle system message instructions: str | None = None chat_template_messages: list[dict[str, Any]] = [] if request.system: if isinstance(request.system, str): instructions = request.system else: instructions = "".join(block.text for block in request.system) instructions = _strip_volatile_headers(instructions) chat_template_messages.append({"role": "system", "content": instructions}) # Convert messages to input input_messages: list[InputMessage] = [] for msg in request.messages: if isinstance(msg.content, str): input_messages.append(InputMessage(role=msg.role, content=msg.content)) chat_template_messages.append({"role": msg.role, "content": msg.content}) continue # Process structured content blocks text_parts: list[str] = [] thinking_parts: list[str] = [] tool_calls: list[dict[str, Any]] = [] tool_results: list[ClaudeToolResultBlock] = [] for block in msg.content: if isinstance(block, ClaudeTextBlock): text_parts.append(block.text) elif isinstance(block, ClaudeThinkingBlock): thinking_parts.append(block.thinking) elif isinstance(block, ClaudeToolUseBlock): tool_calls.append( { "id": block.id, "type": "function", "function": { "name": block.name, "arguments": json.dumps(block.input), }, } ) elif isinstance(block, ClaudeToolResultBlock): tool_results.append(block) content = "".join(text_parts) reasoning_content = "".join(thinking_parts) if thinking_parts else None # Build InputMessage from text content if msg.role in ("user", "assistant"): input_messages.append(InputMessage(role=msg.role, content=content)) # Build chat_template_messages preserving tool structure if tool_calls: chat_msg: dict[str, Any] = { "role": "assistant", "content": content, "tool_calls": tool_calls, } if reasoning_content: chat_msg["reasoning_content"] = reasoning_content chat_template_messages.append(chat_msg) elif tool_results: for tr in tool_results: chat_template_messages.append( { "role": "tool", "tool_call_id": tr.tool_use_id, "content": _extract_tool_result_text(tr), } ) else: chat_msg = {"role": msg.role, "content": content} if reasoning_content: chat_msg["reasoning_content"] = reasoning_content chat_template_messages.append(chat_msg) # Convert Claude tool definitions to OpenAI-style function tools tools: list[dict[str, Any]] | None = None if request.tools: tools = [ { "type": "function", "function": { "name": tool.name, "description": tool.description or "", "parameters": tool.input_schema, }, } for tool in request.tools ] enable_thinking: bool | None = None if request.thinking is not None: enable_thinking = request.thinking.type in ("enabled", "adaptive") return TextGenerationTaskParams( model=request.model, input=input_messages if input_messages else [InputMessage(role="user", content="")], instructions=instructions, max_output_tokens=request.max_tokens, temperature=request.temperature, top_p=request.top_p, top_k=request.top_k, stop=request.stop_sequences, stream=request.stream, tools=tools, enable_thinking=enable_thinking, chat_template_messages=chat_template_messages if chat_template_messages else None, ) async def collect_claude_response( command_id: CommandId, model: str, chunk_stream: AsyncGenerator[ ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None ], ) -> AsyncGenerator[str]: # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason """Collect all token chunks and return a single ClaudeMessagesResponse.""" text_parts: list[str] = [] thinking_parts: list[str] = [] tool_use_blocks: list[ClaudeToolUseBlock] = [] stop_reason: ClaudeStopReason | None = None last_usage: Usage | None = None error_message: str | None = None async for chunk in chunk_stream: if isinstance(chunk, PrefillProgressChunk): continue if isinstance(chunk, ErrorChunk): error_message = chunk.error_message or "Internal server error" break last_usage = chunk.usage or last_usage if isinstance(chunk, ToolCallChunk): for tool in chunk.tool_calls: tool_use_blocks.append( ClaudeToolUseBlock( id=f"toolu_{tool.id}", name=tool.name, input=json.loads(tool.arguments), # pyright: ignore[reportAny] ) ) stop_reason = "tool_use" continue if chunk.is_thinking: thinking_parts.append(chunk.text) else: text_parts.append(chunk.text) if chunk.finish_reason is not None: stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason) if error_message is not None: raise ValueError(error_message) combined_text = "".join(text_parts) combined_thinking = "".join(thinking_parts) # Build content blocks content: list[ClaudeContentBlock] = [] if combined_thinking: content.append(ClaudeThinkingBlock(thinking=combined_thinking)) if combined_text: content.append(ClaudeTextBlock(text=combined_text)) content.extend(tool_use_blocks) # If no content at all, include empty text block if not content: content.append(ClaudeTextBlock(text="")) # Use actual usage data if available input_tokens = last_usage.prompt_tokens if last_usage else 0 output_tokens = last_usage.completion_tokens if last_usage else 0 yield ClaudeMessagesResponse( id=f"msg_{command_id}", model=model, content=content, stop_reason=stop_reason, usage=ClaudeUsage( input_tokens=input_tokens, output_tokens=output_tokens, ), ).model_dump_json() return async def generate_claude_stream( command_id: CommandId, model: str, chunk_stream: AsyncGenerator[ ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None ], ) -> AsyncGenerator[str, None]: """Generate Claude Messages API streaming events from TokenChunks.""" # Initial message_start event initial_message = ClaudeMessageStart( id=f"msg_{command_id}", model=model, content=[], stop_reason=None, usage=ClaudeUsage(input_tokens=0, output_tokens=0), ) start_event = ClaudeMessageStartEvent(message=initial_message) yield f"event: message_start\ndata: {start_event.model_dump_json()}\n\n" output_tokens = 0 stop_reason: ClaudeStopReason | None = None last_usage: Usage | None = None next_block_index = 0 # Track whether we've started thinking/text blocks thinking_block_started = False thinking_block_index = -1 text_block_started = False text_block_index = -1 async for chunk in chunk_stream: if isinstance(chunk, PrefillProgressChunk): continue if isinstance(chunk, ErrorChunk): # Close text block and bail break last_usage = chunk.usage or last_usage if isinstance(chunk, ToolCallChunk): stop_reason = "tool_use" # Emit tool_use content blocks for tool in chunk.tool_calls: tool_id = f"toolu_{tool.id}" tool_input_json = tool.arguments # content_block_start for tool_use tool_block_start = ClaudeContentBlockStartEvent( index=next_block_index, content_block=ClaudeToolUseBlock( id=tool_id, name=tool.name, input={} ), ) yield f"event: content_block_start\ndata: {tool_block_start.model_dump_json()}\n\n" # content_block_delta with input_json_delta tool_delta_event = ClaudeContentBlockDeltaEvent( index=next_block_index, delta=ClaudeInputJsonDelta(partial_json=tool_input_json), ) yield f"event: content_block_delta\ndata: {tool_delta_event.model_dump_json()}\n\n" # content_block_stop tool_block_stop = ClaudeContentBlockStopEvent(index=next_block_index) yield f"event: content_block_stop\ndata: {tool_block_stop.model_dump_json()}\n\n" next_block_index += 1 continue output_tokens += 1 # Count each chunk as one token if chunk.is_thinking: # Start thinking block on first thinking token if not thinking_block_started: thinking_block_started = True thinking_block_index = next_block_index next_block_index += 1 block_start = ClaudeContentBlockStartEvent( index=thinking_block_index, content_block=ClaudeThinkingBlock(thinking=""), ) yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n" delta_event = ClaudeContentBlockDeltaEvent( index=thinking_block_index, delta=ClaudeThinkingDelta(thinking=chunk.text), ) yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n" else: # Close thinking block when transitioning to text if thinking_block_started and text_block_index == -1: block_stop = ClaudeContentBlockStopEvent(index=thinking_block_index) yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n" # Start text block on first text token if not text_block_started: text_block_started = True text_block_index = next_block_index next_block_index += 1 block_start = ClaudeContentBlockStartEvent( index=text_block_index, content_block=ClaudeTextBlock(text=""), ) yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n" delta_event = ClaudeContentBlockDeltaEvent( index=text_block_index, delta=ClaudeTextDelta(text=chunk.text), ) yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n" if chunk.finish_reason is not None: stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason) # Use actual token count from usage if available if last_usage is not None: output_tokens = last_usage.completion_tokens # Close any open blocks if thinking_block_started and text_block_index == -1: block_stop = ClaudeContentBlockStopEvent(index=thinking_block_index) yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n" if text_block_started: block_stop = ClaudeContentBlockStopEvent(index=text_block_index) yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n" if not thinking_block_started and not text_block_started: empty_start = ClaudeContentBlockStartEvent( index=0, content_block=ClaudeTextBlock(text="") ) yield f"event: content_block_start\ndata: {empty_start.model_dump_json()}\n\n" empty_stop = ClaudeContentBlockStopEvent(index=0) yield f"event: content_block_stop\ndata: {empty_stop.model_dump_json()}\n\n" # message_delta message_delta = ClaudeMessageDeltaEvent( delta=ClaudeMessageDelta(stop_reason=stop_reason), usage=ClaudeMessageDeltaUsage(output_tokens=output_tokens), ) yield f"event: message_delta\ndata: {message_delta.model_dump_json()}\n\n" # message_stop message_stop = ClaudeMessageStopEvent() yield f"event: message_stop\ndata: {message_stop.model_dump_json()}\n\n" ================================================ FILE: src/exo/api/adapters/ollama.py ================================================ from __future__ import annotations import json from collections.abc import AsyncGenerator from typing import Any from exo.api.types.ollama_api import ( OllamaChatRequest, OllamaChatResponse, OllamaDoneReason, OllamaGenerateRequest, OllamaGenerateResponse, OllamaMessage, OllamaToolCall, OllamaToolFunction, ) from exo.shared.types.chunks import ( ErrorChunk, PrefillProgressChunk, TokenChunk, ToolCallChunk, ) from exo.shared.types.common import CommandId from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams def _map_done_reason( finish_reason: str | None, ) -> OllamaDoneReason | None: if finish_reason is None: return None if finish_reason == "stop": return "stop" if finish_reason == "length": return "length" if finish_reason in ("tool_calls", "function_call"): return "tool_call" if finish_reason == "error": return "error" return "stop" def _try_parse_json(value: str) -> dict[str, Any] | str: try: return json.loads(value) # type: ignore except json.JSONDecodeError: return value def _build_tool_calls(chunk: ToolCallChunk) -> list[OllamaToolCall]: tool_calls: list[OllamaToolCall] = [] for index, tool in enumerate(chunk.tool_calls): # tool.arguments is always str; try to parse as JSON dict for Ollama format arguments: dict[str, Any] | str = _try_parse_json(tool.arguments) tool_calls.append( OllamaToolCall( id=tool.id, type="function", function=OllamaToolFunction( name=tool.name, arguments=arguments, index=index ), ) ) return tool_calls def _get_usage( chunk: TokenChunk | ToolCallChunk, ) -> tuple[int | None, int | None]: """Extract (prompt_eval_count, eval_count) from a chunk.""" if chunk.usage is not None: return (chunk.usage.prompt_tokens, chunk.usage.completion_tokens) if chunk.stats is not None: return (chunk.stats.prompt_tokens, chunk.stats.generation_tokens) return (None, None) def ollama_request_to_text_generation( request: OllamaChatRequest, ) -> TextGenerationTaskParams: """Convert Ollama chat request to exo's internal text generation format.""" instructions: str | None = None input_messages: list[InputMessage] = [] chat_template_messages: list[dict[str, Any]] = [] tool_message_index = 0 for msg in request.messages: content = msg.content or "" if msg.role == "system": if instructions is None: instructions = content else: instructions = f"{instructions}\n{content}" chat_template_messages.append({"role": "system", "content": content}) continue if msg.role in ("user", "assistant") and ( msg.content is not None or msg.thinking is not None or msg.tool_calls ): input_messages.append(InputMessage(role=msg.role, content=content)) dumped: dict[str, Any] = {"role": msg.role, "content": content} if msg.thinking is not None: dumped["thinking"] = msg.thinking if msg.tool_calls is not None: tool_calls_list: list[dict[str, Any]] = [] for tc in msg.tool_calls: function: dict[str, Any] = { "name": tc.function.name, "arguments": ( json.dumps(tc.function.arguments) if isinstance(tc.function.arguments, dict) else tc.function.arguments ), } if tc.function.index is not None: function["index"] = tc.function.index tool_call: dict[str, Any] = {"function": function} if tc.id is not None: tool_call["id"] = tc.id if tc.type is not None: tool_call["type"] = tc.type tool_calls_list.append(tool_call) dumped["tool_calls"] = tool_calls_list if msg.name is not None: dumped["name"] = msg.name if msg.role == "tool": tool_message_index += 1 tool_call_id = msg.tool_name or msg.name or f"tool_{tool_message_index}" dumped["tool_call_id"] = tool_call_id if msg.tool_name is not None: dumped["tool_name"] = msg.tool_name chat_template_messages.append(dumped) options = request.options return TextGenerationTaskParams( model=request.model, input=input_messages if input_messages else [InputMessage(role="user", content="")], instructions=instructions, max_output_tokens=options.num_predict if options else None, temperature=options.temperature if options else None, top_p=options.top_p if options else None, top_k=options.top_k if options else None, stop=options.stop if options else None, seed=options.seed if options else None, stream=request.stream, tools=request.tools, enable_thinking=request.think, chat_template_messages=chat_template_messages if chat_template_messages else None, ) async def generate_ollama_chat_stream( _command_id: CommandId, chunk_stream: AsyncGenerator[ ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None ], ) -> AsyncGenerator[str, None]: """Generate streaming responses in Ollama format (newline-delimited JSON).""" thinking_parts: list[str] = [] async for chunk in chunk_stream: match chunk: case PrefillProgressChunk(): continue case ErrorChunk(): error_response = OllamaChatResponse( model=str(chunk.model), message=OllamaMessage( role="assistant", content=chunk.error_message ), done=True, done_reason="error", ) yield f"{error_response.model_dump_json(exclude_none=True)}\n" return case ToolCallChunk(): prompt_eval, eval_count = _get_usage(chunk) response = OllamaChatResponse( model=str(chunk.model), message=OllamaMessage( role="assistant", content="", tool_calls=_build_tool_calls(chunk), thinking="".join(thinking_parts) if thinking_parts else None, ), done=True, done_reason="tool_call", prompt_eval_count=prompt_eval, eval_count=eval_count, ) yield f"{response.model_dump_json(exclude_none=True)}\n" return case TokenChunk(): done = chunk.finish_reason is not None if chunk.is_thinking: thinking_parts.append(chunk.text) response = OllamaChatResponse( model=str(chunk.model), message=OllamaMessage( role="assistant", content="", thinking=chunk.text ), done=False, ) yield f"{response.model_dump_json(exclude_none=True)}\n" elif done: prompt_eval, eval_count = _get_usage(chunk) response = OllamaChatResponse( model=str(chunk.model), message=OllamaMessage( role="assistant", content=chunk.text, ), done=True, done_reason=_map_done_reason(chunk.finish_reason), prompt_eval_count=prompt_eval, eval_count=eval_count, ) yield f"{response.model_dump_json(exclude_none=True)}\n" else: response = OllamaChatResponse( model=str(chunk.model), message=OllamaMessage(role="assistant", content=chunk.text), done=False, ) yield f"{response.model_dump_json(exclude_none=True)}\n" if done: return async def collect_ollama_chat_response( _command_id: CommandId, chunk_stream: AsyncGenerator[ ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None ], ) -> AsyncGenerator[str]: """Collect streaming chunks into a single non-streaming Ollama response. Returns an AsyncGenerator[str] (single yield) for consistency with FastAPI StreamingResponse cancellation handling. """ text_parts: list[str] = [] thinking_parts: list[str] = [] tool_calls: list[OllamaToolCall] = [] model: str | None = None finish_reason: str | None = None prompt_eval_count: int | None = None eval_count: int | None = None async for chunk in chunk_stream: match chunk: case PrefillProgressChunk(): continue case ErrorChunk(): raise ValueError(chunk.error_message or "Internal server error") case TokenChunk(): if model is None: model = str(chunk.model) if chunk.is_thinking: thinking_parts.append(chunk.text) else: text_parts.append(chunk.text) if chunk.finish_reason is not None: finish_reason = chunk.finish_reason prompt_eval_count, eval_count = _get_usage(chunk) case ToolCallChunk(): if model is None: model = str(chunk.model) tool_calls.extend(_build_tool_calls(chunk)) finish_reason = chunk.finish_reason prompt_eval_count, eval_count = _get_usage(chunk) combined_text = "".join(text_parts) combined_thinking = "".join(thinking_parts) if thinking_parts else None assert model is not None yield OllamaChatResponse( model=model, message=OllamaMessage( role="assistant", content=combined_text, thinking=combined_thinking, tool_calls=tool_calls if tool_calls else None, ), done=True, done_reason=_map_done_reason(finish_reason), prompt_eval_count=prompt_eval_count, eval_count=eval_count, ).model_dump_json(exclude_none=True) return # ── /api/generate ── def ollama_generate_request_to_text_generation( request: OllamaGenerateRequest, ) -> TextGenerationTaskParams: """Convert Ollama generate request to exo's internal text generation format.""" chat_template_messages: list[dict[str, Any]] = [] if request.system: chat_template_messages.append({"role": "system", "content": request.system}) chat_template_messages.append({"role": "user", "content": request.prompt}) options = request.options return TextGenerationTaskParams( model=request.model, input=[InputMessage(role="user", content=request.prompt)], instructions=request.system, max_output_tokens=options.num_predict if options else None, temperature=options.temperature if options else None, top_p=options.top_p if options else None, top_k=options.top_k if options else None, stop=options.stop if options else None, seed=options.seed if options else None, stream=request.stream, enable_thinking=request.think, chat_template_messages=chat_template_messages if chat_template_messages else None, ) async def generate_ollama_generate_stream( _command_id: CommandId, chunk_stream: AsyncGenerator[ ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None ], ) -> AsyncGenerator[str, None]: """Generate streaming responses for /api/generate in Ollama NDJSON format.""" thinking_parts: list[str] = [] async for chunk in chunk_stream: match chunk: case PrefillProgressChunk(): continue case ErrorChunk(): resp = OllamaGenerateResponse( model=str(chunk.model), response="", done=True, done_reason="error", ) yield f"{resp.model_dump_json(exclude_none=True)}\n" return case ToolCallChunk(): # generate endpoint doesn't support tools; emit as done prompt_eval, eval_count = _get_usage(chunk) resp = OllamaGenerateResponse( model=str(chunk.model), response="", done=True, done_reason="stop", prompt_eval_count=prompt_eval, eval_count=eval_count, ) yield f"{resp.model_dump_json(exclude_none=True)}\n" return case TokenChunk(): done = chunk.finish_reason is not None if chunk.is_thinking: thinking_parts.append(chunk.text) resp = OllamaGenerateResponse( model=str(chunk.model), response="", thinking=chunk.text, done=False, ) yield f"{resp.model_dump_json(exclude_none=True)}\n" elif done: prompt_eval, eval_count = _get_usage(chunk) resp = OllamaGenerateResponse( model=str(chunk.model), response=chunk.text, done=True, done_reason=_map_done_reason(chunk.finish_reason), prompt_eval_count=prompt_eval, eval_count=eval_count, ) yield f"{resp.model_dump_json(exclude_none=True)}\n" else: resp = OllamaGenerateResponse( model=str(chunk.model), response=chunk.text, done=False, ) yield f"{resp.model_dump_json(exclude_none=True)}\n" if done: return async def collect_ollama_generate_response( _command_id: CommandId, chunk_stream: AsyncGenerator[ ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None ], ) -> AsyncGenerator[str]: """Collect chunks into a single non-streaming /api/generate response.""" text_parts: list[str] = [] thinking_parts: list[str] = [] model: str | None = None finish_reason: str | None = None prompt_eval_count: int | None = None eval_count: int | None = None async for chunk in chunk_stream: match chunk: case PrefillProgressChunk(): continue case ErrorChunk(): raise ValueError(chunk.error_message or "Internal server error") case TokenChunk(): if model is None: model = str(chunk.model) if chunk.is_thinking: thinking_parts.append(chunk.text) else: text_parts.append(chunk.text) if chunk.finish_reason is not None: finish_reason = chunk.finish_reason prompt_eval_count, eval_count = _get_usage(chunk) case ToolCallChunk(): if model is None: model = str(chunk.model) finish_reason = chunk.finish_reason prompt_eval_count, eval_count = _get_usage(chunk) assert model is not None yield OllamaGenerateResponse( model=model, response="".join(text_parts), thinking="".join(thinking_parts) if thinking_parts else None, done=True, done_reason=_map_done_reason(finish_reason), prompt_eval_count=prompt_eval_count, eval_count=eval_count, ).model_dump_json(exclude_none=True) return ================================================ FILE: src/exo/api/adapters/responses.py ================================================ """OpenAI Responses API adapter for converting requests/responses.""" from collections.abc import AsyncGenerator from itertools import count from typing import Any from exo.api.types import Usage from exo.api.types.openai_responses import ( FunctionCallInputItem, ResponseCompletedEvent, ResponseContentPart, ResponseContentPartAddedEvent, ResponseContentPartDoneEvent, ResponseCreatedEvent, ResponseFunctionCallArgumentsDeltaEvent, ResponseFunctionCallArgumentsDoneEvent, ResponseFunctionCallItem, ResponseInProgressEvent, ResponseInputMessage, ResponseItem, ResponseMessageItem, ResponseOutputItemAddedEvent, ResponseOutputItemDoneEvent, ResponseOutputText, ResponseReasoningItem, ResponseReasoningSummaryPartAddedEvent, ResponseReasoningSummaryPartDoneEvent, ResponseReasoningSummaryText, ResponseReasoningSummaryTextDeltaEvent, ResponseReasoningSummaryTextDoneEvent, ResponsesRequest, ResponsesResponse, ResponsesStreamEvent, ResponseTextDeltaEvent, ResponseTextDoneEvent, ResponseUsage, ) from exo.shared.types.chunks import ( ErrorChunk, PrefillProgressChunk, TokenChunk, ToolCallChunk, ) from exo.shared.types.common import CommandId from exo.shared.types.text_generation import ( InputMessage, TextGenerationTaskParams, resolve_reasoning_params, ) def _format_sse(event: ResponsesStreamEvent) -> str: """Format a streaming event as an SSE message.""" return f"event: {event.type}\ndata: {event.model_dump_json()}\n\n" def _extract_content(content: str | list[ResponseContentPart]) -> str: """Extract plain text from a content field that may be a string or list of parts.""" if isinstance(content, str): return content return "".join(part.text for part in content) def responses_request_to_text_generation( request: ResponsesRequest, ) -> TextGenerationTaskParams: input_value: list[InputMessage] built_chat_template: list[dict[str, Any]] | None = None if isinstance(request.input, str): input_value = [InputMessage(role="user", content=request.input)] else: input_messages: list[InputMessage] = [] chat_template_messages: list[dict[str, Any]] = [] if request.instructions is not None: chat_template_messages.append( {"role": "system", "content": request.instructions} ) for item in request.input: if isinstance(item, ResponseInputMessage): content = _extract_content(item.content) if item.role in ("user", "assistant", "developer"): input_messages.append(InputMessage(role=item.role, content=content)) if item.role == "system": chat_template_messages.append( {"role": "system", "content": content} ) else: chat_template_messages.append( {"role": item.role, "content": content} ) elif isinstance(item, FunctionCallInputItem): chat_template_messages.append( { "role": "assistant", "content": "", "tool_calls": [ { "id": item.call_id, "type": "function", "function": { "name": item.name, "arguments": item.arguments, }, } ], } ) else: chat_template_messages.append( { "role": "tool", "tool_call_id": item.call_id, "content": item.output, } ) input_value = ( input_messages if input_messages else [InputMessage(role="user", content="")] ) built_chat_template = chat_template_messages if chat_template_messages else None effort_from_reasoning = request.reasoning.effort if request.reasoning else None resolved_effort, resolved_thinking = resolve_reasoning_params( effort_from_reasoning, request.enable_thinking ) # The responses API often does not provide tool args nested under a "function" field. # Since we follow the chat completions format of tools in the backend (for MLX chat templates) # we need to normalise to this format. normalised_tools: list[dict[str, Any]] | None = None if request.tools: normalised_tools = [] for tool in request.tools: if "function" in tool: normalised_tools.append(tool) else: normalised_tools.append( { "type": "function", "function": { "name": tool.get("name", ""), "description": tool.get("description", ""), "parameters": tool.get("parameters", {}), **({"strict": tool["strict"]} if "strict" in tool else {}), }, } ) return TextGenerationTaskParams( model=request.model, input=input_value, instructions=request.instructions, max_output_tokens=request.max_output_tokens, temperature=request.temperature, top_p=request.top_p, stream=request.stream, tools=normalised_tools, top_k=request.top_k, stop=request.stop, seed=request.seed, chat_template_messages=built_chat_template or request.chat_template_messages, reasoning_effort=resolved_effort, enable_thinking=resolved_thinking, ) async def collect_responses_response( command_id: CommandId, model: str, chunk_stream: AsyncGenerator[ ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None ], ) -> AsyncGenerator[str]: # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason """Collect all token chunks and return a single ResponsesResponse.""" response_id = f"resp_{command_id}" item_id = f"item_{command_id}" reasoning_id = f"rs_{command_id}" accumulated_text = "" thinking_parts: list[str] = [] function_call_items: list[ResponseFunctionCallItem] = [] last_usage: Usage | None = None error_message: str | None = None async for chunk in chunk_stream: if isinstance(chunk, PrefillProgressChunk): continue if isinstance(chunk, ErrorChunk): error_message = chunk.error_message or "Internal server error" break last_usage = chunk.usage or last_usage if isinstance(chunk, ToolCallChunk): for tool in chunk.tool_calls: function_call_items.append( ResponseFunctionCallItem( id=tool.id, call_id=tool.id, name=tool.name, arguments=tool.arguments, ) ) continue if chunk.is_thinking: thinking_parts.append(chunk.text) continue accumulated_text += chunk.text if error_message is not None: raise ValueError(error_message) # Create usage from usage data if available usage = None if last_usage is not None: usage = ResponseUsage( input_tokens=last_usage.prompt_tokens, output_tokens=last_usage.completion_tokens, total_tokens=last_usage.total_tokens, ) output: list[ResponseItem] = [] if thinking_parts: output.append( ResponseReasoningItem( id=reasoning_id, summary=[ResponseReasoningSummaryText(text="".join(thinking_parts))], ) ) output.append( ResponseMessageItem( id=item_id, content=[ResponseOutputText(text=accumulated_text)], status="completed", ) ) output.extend(function_call_items) yield ResponsesResponse( id=response_id, model=model, status="completed", output=output, output_text=accumulated_text, usage=usage, ).model_dump_json() return async def generate_responses_stream( command_id: CommandId, model: str, chunk_stream: AsyncGenerator[ ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None ], ) -> AsyncGenerator[str, None]: """Generate OpenAI Responses API streaming events from TokenChunks.""" response_id = f"resp_{command_id}" item_id = f"item_{command_id}" reasoning_id = f"rs_{command_id}" seq = count(1) # response.created initial_response = ResponsesResponse( id=response_id, model=model, status="in_progress", output=[], output_text="", ) created_event = ResponseCreatedEvent( sequence_number=next(seq), response=initial_response ) yield _format_sse(created_event) # response.in_progress in_progress_event = ResponseInProgressEvent( sequence_number=next(seq), response=initial_response ) yield _format_sse(in_progress_event) accumulated_text = "" accumulated_thinking = "" function_call_items: list[ResponseFunctionCallItem] = [] last_usage: Usage | None = None next_output_index = 0 # Track dynamic block creation reasoning_started = False reasoning_output_index = -1 message_started = False message_output_index = -1 async for chunk in chunk_stream: if isinstance(chunk, PrefillProgressChunk): continue if isinstance(chunk, ErrorChunk): break last_usage = chunk.usage or last_usage if isinstance(chunk, ToolCallChunk): for tool in chunk.tool_calls: fc_id = f"fc_{tool.id}" call_id = f"call_{tool.id}" # response.output_item.added for function_call fc_item = ResponseFunctionCallItem( id=fc_id, call_id=call_id, name=tool.name, arguments="", status="in_progress", ) fc_added = ResponseOutputItemAddedEvent( sequence_number=next(seq), output_index=next_output_index, item=fc_item, ) yield _format_sse(fc_added) # response.function_call_arguments.delta args_delta = ResponseFunctionCallArgumentsDeltaEvent( sequence_number=next(seq), item_id=fc_id, output_index=next_output_index, delta=tool.arguments, ) yield _format_sse(args_delta) # response.function_call_arguments.done args_done = ResponseFunctionCallArgumentsDoneEvent( sequence_number=next(seq), item_id=fc_id, output_index=next_output_index, name=tool.name, arguments=tool.arguments, ) yield _format_sse(args_done) # response.output_item.done fc_done_item = ResponseFunctionCallItem( id=fc_id, call_id=call_id, name=tool.name, arguments=tool.arguments, status="completed", ) fc_item_done = ResponseOutputItemDoneEvent( sequence_number=next(seq), output_index=next_output_index, item=fc_done_item, ) yield _format_sse(fc_item_done) function_call_items.append(fc_done_item) next_output_index += 1 continue if chunk.is_thinking: # Start reasoning block on first thinking token if not reasoning_started: reasoning_started = True reasoning_output_index = next_output_index next_output_index += 1 # response.output_item.added for reasoning reasoning_item = ResponseReasoningItem( id=reasoning_id, summary=[], status="in_progress", ) rs_added = ResponseOutputItemAddedEvent( sequence_number=next(seq), output_index=reasoning_output_index, item=reasoning_item, ) yield _format_sse(rs_added) # response.reasoning_summary_part.added part_added = ResponseReasoningSummaryPartAddedEvent( sequence_number=next(seq), item_id=reasoning_id, output_index=reasoning_output_index, summary_index=0, part=ResponseReasoningSummaryText(text=""), ) yield _format_sse(part_added) accumulated_thinking += chunk.text # response.reasoning_summary_text.delta rs_delta = ResponseReasoningSummaryTextDeltaEvent( sequence_number=next(seq), item_id=reasoning_id, output_index=reasoning_output_index, summary_index=0, delta=chunk.text, ) yield _format_sse(rs_delta) continue # Close reasoning block when transitioning to text if reasoning_started and not message_started: # response.reasoning_summary_text.done rs_text_done = ResponseReasoningSummaryTextDoneEvent( sequence_number=next(seq), item_id=reasoning_id, output_index=reasoning_output_index, summary_index=0, text=accumulated_thinking, ) yield _format_sse(rs_text_done) # response.reasoning_summary_part.done rs_part_done = ResponseReasoningSummaryPartDoneEvent( sequence_number=next(seq), item_id=reasoning_id, output_index=reasoning_output_index, summary_index=0, part=ResponseReasoningSummaryText(text=accumulated_thinking), ) yield _format_sse(rs_part_done) # response.output_item.done for reasoning rs_item_done = ResponseOutputItemDoneEvent( sequence_number=next(seq), output_index=reasoning_output_index, item=ResponseReasoningItem( id=reasoning_id, summary=[ResponseReasoningSummaryText(text=accumulated_thinking)], ), ) yield _format_sse(rs_item_done) # Start message block on first text token if not message_started: message_started = True message_output_index = next_output_index next_output_index += 1 initial_item = ResponseMessageItem( id=item_id, content=[ResponseOutputText(text="")], status="in_progress", ) item_added = ResponseOutputItemAddedEvent( sequence_number=next(seq), output_index=message_output_index, item=initial_item, ) yield _format_sse(item_added) initial_part = ResponseOutputText(text="") part_added = ResponseContentPartAddedEvent( sequence_number=next(seq), item_id=item_id, output_index=message_output_index, content_index=0, part=initial_part, ) yield _format_sse(part_added) accumulated_text += chunk.text # response.output_text.delta delta_event = ResponseTextDeltaEvent( sequence_number=next(seq), item_id=item_id, output_index=message_output_index, content_index=0, delta=chunk.text, ) yield _format_sse(delta_event) # Close reasoning block if it was never followed by text if reasoning_started and not message_started: rs_text_done = ResponseReasoningSummaryTextDoneEvent( sequence_number=next(seq), item_id=reasoning_id, output_index=reasoning_output_index, summary_index=0, text=accumulated_thinking, ) yield _format_sse(rs_text_done) rs_part_done = ResponseReasoningSummaryPartDoneEvent( sequence_number=next(seq), item_id=reasoning_id, output_index=reasoning_output_index, summary_index=0, part=ResponseReasoningSummaryText(text=accumulated_thinking), ) yield _format_sse(rs_part_done) rs_item_done = ResponseOutputItemDoneEvent( sequence_number=next(seq), output_index=reasoning_output_index, item=ResponseReasoningItem( id=reasoning_id, summary=[ResponseReasoningSummaryText(text=accumulated_thinking)], ), ) yield _format_sse(rs_item_done) # If no message block was started, create one now (empty text) if not message_started: message_output_index = next_output_index next_output_index += 1 initial_item = ResponseMessageItem( id=item_id, content=[ResponseOutputText(text="")], status="in_progress", ) item_added = ResponseOutputItemAddedEvent( sequence_number=next(seq), output_index=message_output_index, item=initial_item, ) yield _format_sse(item_added) initial_part = ResponseOutputText(text="") part_added_evt = ResponseContentPartAddedEvent( sequence_number=next(seq), item_id=item_id, output_index=message_output_index, content_index=0, part=initial_part, ) yield _format_sse(part_added_evt) # response.output_text.done text_done = ResponseTextDoneEvent( sequence_number=next(seq), item_id=item_id, output_index=message_output_index, content_index=0, text=accumulated_text, ) yield _format_sse(text_done) # response.content_part.done final_part = ResponseOutputText(text=accumulated_text) part_done = ResponseContentPartDoneEvent( sequence_number=next(seq), item_id=item_id, output_index=message_output_index, content_index=0, part=final_part, ) yield _format_sse(part_done) # response.output_item.done final_message_item = ResponseMessageItem( id=item_id, content=[ResponseOutputText(text=accumulated_text)], status="completed", ) item_done = ResponseOutputItemDoneEvent( sequence_number=next(seq), output_index=message_output_index, item=final_message_item, ) yield _format_sse(item_done) # Create usage from usage data if available usage = None if last_usage is not None: usage = ResponseUsage( input_tokens=last_usage.prompt_tokens, output_tokens=last_usage.completion_tokens, total_tokens=last_usage.total_tokens, ) # response.completed output: list[ResponseItem] = [] if reasoning_started: output.append( ResponseReasoningItem( id=reasoning_id, summary=[ResponseReasoningSummaryText(text=accumulated_thinking)], ) ) output.append(final_message_item) output.extend(function_call_items) final_response = ResponsesResponse( id=response_id, model=model, status="completed", output=output, output_text=accumulated_text, usage=usage, ) completed_event = ResponseCompletedEvent( sequence_number=next(seq), response=final_response ) yield _format_sse(completed_event) ================================================ FILE: src/exo/api/main.py ================================================ import base64 import contextlib import json import random import time from collections.abc import AsyncGenerator, Awaitable, Callable, Iterable from datetime import datetime, timezone from http import HTTPStatus from pathlib import Path from typing import Annotated, Literal, cast from uuid import uuid4 import anyio from anyio import BrokenResourceError from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import FileResponse, JSONResponse, StreamingResponse from fastapi.staticfiles import StaticFiles from hypercorn.asyncio import serve # pyright: ignore[reportUnknownVariableType] from hypercorn.config import Config from hypercorn.typing import ASGIFramework from loguru import logger from exo.api.adapters.chat_completions import ( chat_request_to_text_generation, collect_chat_response, generate_chat_stream, ) from exo.api.adapters.claude import ( claude_request_to_text_generation, collect_claude_response, generate_claude_stream, ) from exo.api.adapters.ollama import ( collect_ollama_chat_response, collect_ollama_generate_response, generate_ollama_chat_stream, generate_ollama_generate_stream, ollama_generate_request_to_text_generation, ollama_request_to_text_generation, ) from exo.api.adapters.responses import ( collect_responses_response, generate_responses_stream, responses_request_to_text_generation, ) from exo.api.types import ( AddCustomModelParams, AdvancedImageParams, BenchChatCompletionRequest, BenchChatCompletionResponse, BenchImageGenerationResponse, BenchImageGenerationTaskParams, CancelCommandResponse, ChatCompletionChoice, ChatCompletionMessage, ChatCompletionRequest, ChatCompletionResponse, CreateInstanceParams, CreateInstanceResponse, DeleteDownloadResponse, DeleteInstanceResponse, DeleteTracesRequest, DeleteTracesResponse, ErrorInfo, ErrorResponse, FinishReason, GenerationStats, HuggingFaceSearchResult, ImageData, ImageEditsTaskParams, ImageGenerationResponse, ImageGenerationStats, ImageGenerationTaskParams, ImageListItem, ImageListResponse, ImageSize, ModelList, ModelListModel, PlaceInstanceParams, PlacementPreview, PlacementPreviewResponse, StartDownloadParams, StartDownloadResponse, ToolCall, TraceCategoryStats, TraceEventResponse, TraceListItem, TraceListResponse, TraceRankStats, TraceResponse, TraceStatsResponse, normalize_image_size, ) from exo.api.types.claude_api import ( ClaudeMessagesRequest, ClaudeMessagesResponse, ) from exo.api.types.ollama_api import ( OllamaChatRequest, OllamaChatResponse, OllamaGenerateRequest, OllamaGenerateResponse, OllamaModelDetails, OllamaModelTag, OllamaPsModel, OllamaPsResponse, OllamaShowRequest, OllamaShowResponse, OllamaTagsResponse, ) from exo.api.types.openai_responses import ( ResponsesRequest, ResponsesResponse, ) from exo.master.image_store import ImageStore from exo.master.placement import place_instance as get_instance_placements from exo.shared.apply import apply from exo.shared.constants import ( DASHBOARD_DIR, EXO_CACHE_HOME, EXO_EVENT_LOG_DIR, EXO_IMAGE_CACHE_DIR, EXO_MAX_CHUNK_SIZE, EXO_TRACING_CACHE_DIR, ) from exo.shared.election import ElectionMessage from exo.shared.logging import InterceptLogger from exo.shared.models.model_cards import ( ModelCard, ModelId, delete_custom_card, get_model_cards, is_custom_card, ) from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file from exo.shared.types.chunks import ( ErrorChunk, ImageChunk, InputImageChunk, PrefillProgressChunk, TokenChunk, ToolCallChunk, ) from exo.shared.types.commands import ( Command, CreateInstance, DeleteDownload, DeleteInstance, DownloadCommand, ForwarderCommand, ForwarderDownloadCommand, ImageEdits, ImageGeneration, PlaceInstance, SendInputChunk, StartDownload, TaskCancelled, TaskFinished, TextGeneration, ) from exo.shared.types.common import CommandId, Id, NodeId, SystemId from exo.shared.types.events import ( ChunkGenerated, Event, IndexedEvent, TracesMerged, ) from exo.shared.types.memory import Memory from exo.shared.types.state import State from exo.shared.types.worker.downloads import DownloadCompleted from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta from exo.shared.types.worker.shards import Sharding from exo.utils.banner import print_startup_banner from exo.utils.channels import Receiver, Sender, channel from exo.utils.disk_event_log import DiskEventLog from exo.utils.power_sampler import PowerSampler from exo.utils.task_group import TaskGroup _API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api" ONBOARDING_COMPLETE_FILE = EXO_CACHE_HOME / "onboarding_complete" def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None) -> str: return f"image/{image_format or 'png'}" def _ensure_seed(params: AdvancedImageParams | None) -> AdvancedImageParams: """Ensure advanced params has a seed set for distributed consistency.""" if params is None: return AdvancedImageParams(seed=random.randint(0, 2**32 - 1)) if params.seed is None: return params.model_copy(update={"seed": random.randint(0, 2**32 - 1)}) return params class API: def __init__( self, node_id: NodeId, *, port: int, event_receiver: Receiver[IndexedEvent], command_sender: Sender[ForwarderCommand], download_command_sender: Sender[ForwarderDownloadCommand], # This lets us pause the API if an election is running election_receiver: Receiver[ElectionMessage], ) -> None: self.state = State() self._event_log = DiskEventLog(_API_EVENT_LOG_DIR) self._system_id = SystemId() self.command_sender = command_sender self.download_command_sender = download_command_sender self.event_receiver = event_receiver self.election_receiver = election_receiver self.node_id: NodeId = node_id self.last_completed_election: int = 0 self.port = port self.paused: bool = False self.paused_ev: anyio.Event = anyio.Event() self.app = FastAPI() @self.app.middleware("http") async def _log_requests( # pyright: ignore[reportUnusedFunction] request: Request, call_next: Callable[[Request], Awaitable[StreamingResponse]], ) -> StreamingResponse: logger.debug(f"API request: {request.method} {request.url.path}") return await call_next(request) self._setup_exception_handlers() self._setup_cors() self._setup_routes() self.app.mount( "/", StaticFiles( directory=DASHBOARD_DIR, html=True, ), name="dashboard", ) self._text_generation_queues: dict[ CommandId, Sender[TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk], ] = {} self._image_generation_queues: dict[ CommandId, Sender[ImageChunk | ErrorChunk] ] = {} self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR) self._tg: TaskGroup = TaskGroup() def reset(self, result_clock: int, event_receiver: Receiver[IndexedEvent]): logger.info("Resetting API State") self._event_log.close() self._event_log = DiskEventLog(_API_EVENT_LOG_DIR) self.state = State() self._system_id = SystemId() self._text_generation_queues = {} self._image_generation_queues = {} self.unpause(result_clock) self.event_receiver.close() self.event_receiver = event_receiver self._tg.start_soon(self._apply_state) def unpause(self, result_clock: int): logger.info("Unpausing API") self.last_completed_election = result_clock self.paused = False self.paused_ev.set() self.paused_ev = anyio.Event() def _setup_exception_handlers(self) -> None: self.app.exception_handler(HTTPException)(self.http_exception_handler) async def http_exception_handler( self, _: Request, exc: HTTPException ) -> JSONResponse: err = ErrorResponse( error=ErrorInfo( message=exc.detail, type=HTTPStatus(exc.status_code).phrase, code=exc.status_code, ) ) return JSONResponse(err.model_dump(), status_code=exc.status_code) def _setup_cors(self) -> None: self.app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) def _setup_routes(self) -> None: self.app.get("/node_id")(lambda: self.node_id) self.app.post("/instance")(self.create_instance) self.app.post("/place_instance")(self.place_instance) self.app.get("/instance/placement")(self.get_placement) self.app.get("/instance/previews")(self.get_placement_previews) self.app.get("/instance/{instance_id}")(self.get_instance) self.app.delete("/instance/{instance_id}")(self.delete_instance) self.app.get("/models")(self.get_models) self.app.get("/v1/models")(self.get_models) self.app.post("/models/add")(self.add_custom_model) self.app.delete("/models/custom/{model_id:path}")(self.delete_custom_model) self.app.get("/models/search")(self.search_models) self.app.post("/v1/chat/completions", response_model=None)( self.chat_completions ) self.app.post("/bench/chat/completions")(self.bench_chat_completions) self.app.post("/v1/images/generations", response_model=None)( self.image_generations ) self.app.post("/bench/images/generations")(self.bench_image_generations) self.app.post("/v1/images/edits", response_model=None)(self.image_edits) self.app.post("/bench/images/edits")(self.bench_image_edits) self.app.get("/images")(self.list_images) self.app.get("/images/{image_id}")(self.get_image) self.app.post("/v1/messages", response_model=None)(self.claude_messages) self.app.post("/v1/responses", response_model=None)(self.openai_responses) self.app.post("/v1/cancel/{command_id}")(self.cancel_command) # Ollama API self.app.head("/ollama/")(self.ollama_version) self.app.head("/ollama/api/version")(self.ollama_version) self.app.post("/ollama/api/chat", response_model=None)(self.ollama_chat) self.app.post("/ollama/api/api/chat", response_model=None)(self.ollama_chat) self.app.post("/ollama/api/v1/chat", response_model=None)(self.ollama_chat) self.app.post("/ollama/api/generate", response_model=None)(self.ollama_generate) self.app.get("/ollama/api/tags")(self.ollama_tags) self.app.get("/ollama/api/api/tags")(self.ollama_tags) self.app.get("/ollama/api/v1/tags")(self.ollama_tags) self.app.post("/ollama/api/show")(self.ollama_show) self.app.get("/ollama/api/ps")(self.ollama_ps) self.app.get("/ollama/api/version")(self.ollama_version) self.app.get("/state")(lambda: self.state) self.app.get("/events")(self.stream_events) self.app.post("/download/start")(self.start_download) self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download) self.app.get("/v1/traces")(self.list_traces) self.app.post("/v1/traces/delete")(self.delete_traces) self.app.get("/v1/traces/{task_id}")(self.get_trace) self.app.get("/v1/traces/{task_id}/stats")(self.get_trace_stats) self.app.get("/v1/traces/{task_id}/raw")(self.get_trace_raw) self.app.get("/onboarding")(self.get_onboarding) self.app.post("/onboarding")(self.complete_onboarding) async def place_instance(self, payload: PlaceInstanceParams): command = PlaceInstance( model_card=await ModelCard.load(payload.model_id), sharding=payload.sharding, instance_meta=payload.instance_meta, min_nodes=payload.min_nodes, ) await self._send(command) return CreateInstanceResponse( message="Command received.", command_id=command.command_id, model_card=command.model_card, ) async def create_instance( self, payload: CreateInstanceParams ) -> CreateInstanceResponse: instance = payload.instance model_card = await ModelCard.load(instance.shard_assignments.model_id) required_memory = model_card.storage_size available_memory = self._calculate_total_available_memory() if required_memory > available_memory: raise HTTPException( status_code=400, detail=f"Insufficient memory to create instance. Required: {required_memory.in_gb:.1f}GB, Available: {available_memory.in_gb:.1f}GB", ) command = CreateInstance( instance=instance, ) await self._send(command) return CreateInstanceResponse( message="Command received.", command_id=command.command_id, model_card=model_card, ) async def get_placement( self, model_id: ModelId, sharding: Sharding = Sharding.Pipeline, instance_meta: InstanceMeta = InstanceMeta.MlxRing, min_nodes: int = 1, ) -> Instance: model_card = await ModelCard.load(model_id) try: placements = get_instance_placements( PlaceInstance( model_card=model_card, sharding=sharding, instance_meta=instance_meta, min_nodes=min_nodes, ), node_memory=self.state.node_memory, node_network=self.state.node_network, topology=self.state.topology, current_instances=self.state.instances, ) except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc current_ids = set(self.state.instances.keys()) new_ids = [ instance_id for instance_id in placements if instance_id not in current_ids ] if len(new_ids) != 1: raise HTTPException( status_code=500, detail="Expected exactly one new instance from placement", ) return placements[new_ids[0]] async def get_placement_previews( self, model_id: ModelId, node_ids: Annotated[list[NodeId] | None, Query()] = None, ) -> PlacementPreviewResponse: seen: set[tuple[ModelId, Sharding, InstanceMeta, int]] = set() previews: list[PlacementPreview] = [] required_nodes = set(node_ids) if node_ids else None if len(list(self.state.topology.list_nodes())) == 0: return PlacementPreviewResponse(previews=[]) try: model_card = await ModelCard.load(model_id) except Exception as exc: raise HTTPException( status_code=400, detail=f"Failed to load model card: {exc}" ) from exc instance_combinations: list[tuple[Sharding, InstanceMeta, int]] = [] for sharding in (Sharding.Pipeline, Sharding.Tensor): for instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl): instance_combinations.extend( [ (sharding, instance_meta, i) for i in range( 1, len(list(self.state.topology.list_nodes())) + 1 ) ] ) # TODO: PDD # instance_combinations.append((Sharding.PrefillDecodeDisaggregation, InstanceMeta.MlxRing, 1)) for sharding, instance_meta, min_nodes in instance_combinations: try: placements = get_instance_placements( PlaceInstance( model_card=model_card, sharding=sharding, instance_meta=instance_meta, min_nodes=min_nodes, ), node_memory=self.state.node_memory, node_network=self.state.node_network, topology=self.state.topology, current_instances=self.state.instances, required_nodes=required_nodes, ) except ValueError as exc: if (model_card.model_id, sharding, instance_meta, 0) not in seen: previews.append( PlacementPreview( model_id=model_card.model_id, sharding=sharding, instance_meta=instance_meta, instance=None, error=str(exc), ) ) seen.add((model_card.model_id, sharding, instance_meta, 0)) continue current_ids = set(self.state.instances.keys()) new_instances = [ instance for instance_id, instance in placements.items() if instance_id not in current_ids ] if len(new_instances) != 1: if (model_card.model_id, sharding, instance_meta, 0) not in seen: previews.append( PlacementPreview( model_id=model_card.model_id, sharding=sharding, instance_meta=instance_meta, instance=None, error="Expected exactly one new instance from placement", ) ) seen.add((model_card.model_id, sharding, instance_meta, 0)) continue instance = new_instances[0] shard_assignments = instance.shard_assignments placement_node_ids = list(shard_assignments.node_to_runner.keys()) memory_delta_by_node: dict[str, int] = {} if placement_node_ids: total_bytes = model_card.storage_size.in_bytes per_node = total_bytes // len(placement_node_ids) remainder = total_bytes % len(placement_node_ids) for index, node_id in enumerate(sorted(placement_node_ids, key=str)): extra = 1 if index < remainder else 0 memory_delta_by_node[str(node_id)] = per_node + extra if ( model_card.model_id, sharding, instance_meta, len(placement_node_ids), ) not in seen: previews.append( PlacementPreview( model_id=model_card.model_id, sharding=sharding, instance_meta=instance_meta, instance=instance, memory_delta_by_node=memory_delta_by_node or None, error=None, ) ) seen.add( ( model_card.model_id, sharding, instance_meta, len(placement_node_ids), ) ) return PlacementPreviewResponse(previews=previews) def get_instance(self, instance_id: InstanceId) -> Instance: if instance_id not in self.state.instances: raise HTTPException(status_code=404, detail="Instance not found") return self.state.instances[instance_id] async def delete_instance(self, instance_id: InstanceId) -> DeleteInstanceResponse: if instance_id not in self.state.instances: raise HTTPException(status_code=404, detail="Instance not found") command = DeleteInstance( instance_id=instance_id, ) await self._send(command) return DeleteInstanceResponse( message="Command received.", command_id=command.command_id, instance_id=instance_id, ) async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse: """Cancel an active command by closing its stream and notifying workers.""" sender = self._text_generation_queues.get( command_id ) or self._image_generation_queues.get(command_id) if sender is None: raise HTTPException( status_code=404, detail="Command not found or already completed", ) await self._send(TaskCancelled(cancelled_command_id=command_id)) sender.close() return CancelCommandResponse( message="Command cancelled.", command_id=command_id, ) async def _token_chunk_stream( self, command_id: CommandId ) -> AsyncGenerator[ TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk, None ]: """Yield chunks for a given command until completion. This is the internal low-level stream used by all API adapters. """ try: self._text_generation_queues[command_id], recv = channel[ TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk ]() with recv as token_chunks: async for chunk in token_chunks: yield chunk if isinstance(chunk, PrefillProgressChunk): continue if chunk.finish_reason is not None: break except anyio.get_cancelled_exc_class(): command = TaskCancelled(cancelled_command_id=command_id) with anyio.CancelScope(shield=True): await self.command_sender.send( ForwarderCommand(origin=self._system_id, command=command) ) raise finally: await self._send(TaskFinished(finished_command_id=command_id)) if command_id in self._text_generation_queues: del self._text_generation_queues[command_id] async def _collect_text_generation_with_stats( self, command_id: CommandId ) -> BenchChatCompletionResponse: sampler = PowerSampler(get_node_system=lambda: self.state.node_system) text_parts: list[str] = [] tool_calls: list[ToolCall] = [] model: ModelId | None = None finish_reason: FinishReason | None = None stats: GenerationStats | None = None async with anyio.create_task_group() as tg: tg.start_soon(sampler.run) async for chunk in self._token_chunk_stream(command_id): if isinstance(chunk, PrefillProgressChunk): continue if chunk.finish_reason == "error": raise HTTPException( status_code=500, detail=chunk.error_message or "Internal server error", ) if model is None: model = chunk.model if isinstance(chunk, TokenChunk): text_parts.append(chunk.text) if isinstance(chunk, ToolCallChunk): tool_calls.extend( ToolCall( id=str(uuid4()), index=i, function=tool, ) for i, tool in enumerate(chunk.tool_calls) ) stats = chunk.stats or stats if chunk.finish_reason is not None: finish_reason = chunk.finish_reason tg.cancel_scope.cancel() combined_text = "".join(text_parts) assert model is not None return BenchChatCompletionResponse( id=command_id, created=int(time.time()), model=model, choices=[ ChatCompletionChoice( index=0, message=ChatCompletionMessage( role="assistant", content=combined_text, tool_calls=tool_calls if tool_calls else None, ), finish_reason=finish_reason, ) ], generation_stats=stats, power_usage=sampler.result(), ) async def _trigger_notify_user_to_download_model(self, model_id: ModelId) -> None: logger.warning( "TODO: we should send a notification to the user to download the model" ) async def chat_completions( self, payload: ChatCompletionRequest ) -> ChatCompletionResponse | StreamingResponse: """OpenAI Chat Completions API - adapter.""" task_params = chat_request_to_text_generation(payload) resolved_model = await self._resolve_and_validate_text_model( ModelId(task_params.model) ) task_params = task_params.model_copy(update={"model": resolved_model}) command = TextGeneration(task_params=task_params) await self._send(command) if payload.stream: return StreamingResponse( generate_chat_stream( command.command_id, self._token_chunk_stream(command.command_id), ), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "close", "X-Accel-Buffering": "no", }, ) else: return StreamingResponse( collect_chat_response( command.command_id, self._token_chunk_stream(command.command_id), ), media_type="application/json", ) async def bench_chat_completions( self, payload: BenchChatCompletionRequest ) -> BenchChatCompletionResponse: task_params = chat_request_to_text_generation(payload) resolved_model = await self._resolve_and_validate_text_model( ModelId(task_params.model) ) task_params = task_params.model_copy(update={"model": resolved_model}) task_params = task_params.model_copy(update={"stream": False, "bench": True}) command = TextGeneration(task_params=task_params) await self._send(command) return await self._collect_text_generation_with_stats(command.command_id) async def _resolve_and_validate_text_model(self, model_id: ModelId) -> ModelId: """Validate a text model exists and return the resolved model ID. Raises HTTPException 404 if no instance is found for the model. """ if not any( instance.shard_assignments.model_id == model_id for instance in self.state.instances.values() ): await self._trigger_notify_user_to_download_model(model_id) raise HTTPException( status_code=404, detail=f"No instance found for model {model_id}", ) return model_id async def _validate_image_model(self, model: ModelId) -> ModelId: """Validate model exists and return resolved model ID. Raises HTTPException 404 if no instance is found for the model. """ model_card = await ModelCard.load(model) resolved_model = model_card.model_id if not any( instance.shard_assignments.model_id == resolved_model for instance in self.state.instances.values() ): await self._trigger_notify_user_to_download_model(resolved_model) raise HTTPException( status_code=404, detail=f"No instance found for model {resolved_model}" ) return resolved_model def stream_events(self) -> StreamingResponse: def _generate_json_array(events: Iterable[Event]) -> Iterable[str]: yield "[" first = True for event in events: if not first: yield "," first = False yield event.model_dump_json() yield "]" return StreamingResponse( _generate_json_array(self._event_log.read_all()), media_type="application/json", ) async def get_image(self, image_id: str) -> FileResponse: stored = self._image_store.get(Id(image_id)) if stored is None: raise HTTPException(status_code=404, detail="Image not found or expired") return FileResponse(path=stored.file_path, media_type=stored.content_type) async def list_images(self, request: Request) -> ImageListResponse: """List all stored images.""" stored_images = self._image_store.list_images() return ImageListResponse( data=[ ImageListItem( image_id=img.image_id, url=self._build_image_url(request, img.image_id), content_type=img.content_type, expires_at=img.expires_at, ) for img in stored_images ] ) def _build_image_url(self, request: Request, image_id: Id) -> str: host = request.headers.get("host", f"localhost:{self.port}") scheme = "https" if request.url.scheme == "https" else "http" return f"{scheme}://{host}/v1/images/{image_id}" async def image_generations( self, request: Request, payload: ImageGenerationTaskParams ) -> ImageGenerationResponse | StreamingResponse: """Handle image generation requests. When stream=True and partial_images > 0, returns a StreamingResponse with SSE-formatted events for partial and final images. """ payload = payload.model_copy( update={ "model": await self._validate_image_model(ModelId(payload.model)), "advanced_params": _ensure_seed(payload.advanced_params), } ) command = ImageGeneration( task_params=payload, ) await self._send(command) # Check if streaming is requested if payload.stream and payload.partial_images and payload.partial_images > 0: return StreamingResponse( self._generate_image_stream( request=request, command_id=command.command_id, num_images=payload.n or 1, response_format=payload.response_format or "b64_json", ), media_type="text/event-stream", ) # Non-streaming: collect all image chunks return await self._collect_image_generation( request=request, command_id=command.command_id, num_images=payload.n or 1, response_format=payload.response_format or "b64_json", ) async def _generate_image_stream( self, request: Request, command_id: CommandId, num_images: int, response_format: str, ) -> AsyncGenerator[str, None]: """Generate SSE stream of partial and final images.""" # Track chunks: {(image_index, is_partial): {chunk_index: data}} image_chunks: dict[tuple[int, bool], dict[int, str]] = {} image_total_chunks: dict[tuple[int, bool], int] = {} image_metadata: dict[tuple[int, bool], tuple[int | None, int | None]] = {} images_complete = 0 try: self._image_generation_queues[command_id], recv = channel[ ImageChunk | ErrorChunk ]() with recv as chunks: async for chunk in chunks: if chunk.finish_reason == "error": error_response = ErrorResponse( error=ErrorInfo( message=chunk.error_message or "Internal server error", type="InternalServerError", code=500, ) ) yield f"data: {error_response.model_dump_json()}\n\n" yield "data: [DONE]\n\n" return key = (chunk.image_index, chunk.is_partial) if key not in image_chunks: image_chunks[key] = {} image_total_chunks[key] = chunk.total_chunks image_metadata[key] = ( chunk.partial_index, chunk.total_partials, ) image_chunks[key][chunk.chunk_index] = chunk.data # Check if this image is complete if len(image_chunks[key]) == image_total_chunks[key]: full_data = "".join( image_chunks[key][i] for i in range(len(image_chunks[key])) ) partial_idx, total_partials = image_metadata[key] if chunk.is_partial: # Yield partial image event (always use b64_json for partials) event_data = { "type": "partial", "image_index": chunk.image_index, "partial_index": partial_idx, "total_partials": total_partials, "format": str(chunk.format), "data": { "b64_json": full_data if response_format == "b64_json" else None, }, } yield f"data: {json.dumps(event_data)}\n\n" else: # Final image if response_format == "url": image_bytes = base64.b64decode(full_data) content_type = _format_to_content_type(chunk.format) stored = self._image_store.store( image_bytes, content_type ) url = self._build_image_url(request, stored.image_id) event_data = { "type": "final", "image_index": chunk.image_index, "format": str(chunk.format), "data": {"url": url}, } else: event_data = { "type": "final", "image_index": chunk.image_index, "format": str(chunk.format), "data": {"b64_json": full_data}, } yield f"data: {json.dumps(event_data)}\n\n" images_complete += 1 if images_complete >= num_images: yield "data: [DONE]\n\n" break # Clean up completed image chunks del image_chunks[key] del image_total_chunks[key] del image_metadata[key] except anyio.get_cancelled_exc_class(): command = TaskCancelled(cancelled_command_id=command_id) with anyio.CancelScope(shield=True): await self.command_sender.send( ForwarderCommand(origin=self._system_id, command=command) ) raise finally: await self._send(TaskFinished(finished_command_id=command_id)) if command_id in self._image_generation_queues: del self._image_generation_queues[command_id] async def _collect_image_chunks( self, request: Request | None, command_id: CommandId, num_images: int, response_format: str, capture_stats: bool = False, ) -> tuple[list[ImageData], ImageGenerationStats | None]: """Collect image chunks and optionally capture stats.""" # Track chunks per image: {image_index: {chunk_index: data}} # Only track non-partial (final) images image_chunks: dict[int, dict[int, str]] = {} image_total_chunks: dict[int, int] = {} image_formats: dict[int, Literal["png", "jpeg", "webp"] | None] = {} images_complete = 0 stats: ImageGenerationStats | None = None try: self._image_generation_queues[command_id], recv = channel[ ImageChunk | ErrorChunk ]() while images_complete < num_images: with recv as chunks: async for chunk in chunks: if chunk.finish_reason == "error": raise HTTPException( status_code=500, detail=chunk.error_message or "Internal server error", ) if chunk.is_partial: continue if chunk.image_index not in image_chunks: image_chunks[chunk.image_index] = {} image_total_chunks[chunk.image_index] = chunk.total_chunks image_formats[chunk.image_index] = chunk.format image_chunks[chunk.image_index][chunk.chunk_index] = chunk.data if capture_stats and chunk.stats is not None: stats = chunk.stats if ( len(image_chunks[chunk.image_index]) == image_total_chunks[chunk.image_index] ): images_complete += 1 if images_complete >= num_images: break images: list[ImageData] = [] for image_idx in range(num_images): chunks_dict = image_chunks[image_idx] full_data = "".join(chunks_dict[i] for i in range(len(chunks_dict))) if response_format == "url" and request is not None: image_bytes = base64.b64decode(full_data) content_type = _format_to_content_type(image_formats.get(image_idx)) stored = self._image_store.store(image_bytes, content_type) url = self._build_image_url(request, stored.image_id) images.append(ImageData(b64_json=None, url=url)) else: images.append( ImageData( b64_json=full_data if response_format == "b64_json" else None, url=None, ) ) return (images, stats if capture_stats else None) except anyio.get_cancelled_exc_class(): command = TaskCancelled(cancelled_command_id=command_id) with anyio.CancelScope(shield=True): await self.command_sender.send( ForwarderCommand(origin=self._system_id, command=command) ) raise finally: await self._send(TaskFinished(finished_command_id=command_id)) if command_id in self._image_generation_queues: del self._image_generation_queues[command_id] async def _collect_image_generation( self, request: Request, command_id: CommandId, num_images: int, response_format: str, ) -> ImageGenerationResponse: """Collect all image chunks (non-streaming) and return a single response.""" images, _ = await self._collect_image_chunks( request, command_id, num_images, response_format, capture_stats=False ) return ImageGenerationResponse(data=images) async def _collect_image_generation_with_stats( self, request: Request | None, command_id: CommandId, num_images: int, response_format: str, ) -> BenchImageGenerationResponse: sampler = PowerSampler(get_node_system=lambda: self.state.node_system) images: list[ImageData] = [] stats: ImageGenerationStats | None = None async with anyio.create_task_group() as tg: tg.start_soon(sampler.run) images, stats = await self._collect_image_chunks( request, command_id, num_images, response_format, capture_stats=True ) tg.cancel_scope.cancel() return BenchImageGenerationResponse( data=images, generation_stats=stats, power_usage=sampler.result() ) async def bench_image_generations( self, request: Request, payload: BenchImageGenerationTaskParams ) -> BenchImageGenerationResponse: payload = payload.model_copy( update={ "model": await self._validate_image_model(ModelId(payload.model)), "stream": False, "partial_images": 0, "advanced_params": _ensure_seed(payload.advanced_params), } ) command = ImageGeneration( task_params=payload, ) await self._send(command) return await self._collect_image_generation_with_stats( request=request, command_id=command.command_id, num_images=payload.n or 1, response_format=payload.response_format or "b64_json", ) async def _send_image_edits_command( self, image: UploadFile, prompt: str, model: ModelId, n: int, size: ImageSize, response_format: Literal["url", "b64_json"], input_fidelity: Literal["low", "high"], stream: bool, partial_images: int, bench: bool, quality: Literal["high", "medium", "low"], output_format: Literal["png", "jpeg", "webp"], advanced_params: AdvancedImageParams | None, ) -> ImageEdits: """Prepare and send an image edits command with chunked image upload.""" resolved_model = await self._validate_image_model(model) advanced_params = _ensure_seed(advanced_params) image_content = await image.read() image_data = base64.b64encode(image_content).decode("utf-8") image_strength = 0.7 if input_fidelity == "high" else 0.3 data_chunks = [ image_data[i : i + EXO_MAX_CHUNK_SIZE] for i in range(0, len(image_data), EXO_MAX_CHUNK_SIZE) ] total_chunks = len(data_chunks) command = ImageEdits( task_params=ImageEditsTaskParams( image_data="", total_input_chunks=total_chunks, prompt=prompt, model=resolved_model, n=n, size=size, response_format=response_format, image_strength=image_strength, stream=stream, partial_images=partial_images, bench=bench, quality=quality, output_format=output_format, advanced_params=advanced_params, ), ) logger.info( f"Sending input image: {len(image_data)} bytes in {total_chunks} chunks" ) for chunk_index, chunk_data in enumerate(data_chunks): await self._send( SendInputChunk( chunk=InputImageChunk( model=resolved_model, command_id=command.command_id, data=chunk_data, chunk_index=chunk_index, total_chunks=total_chunks, ) ) ) await self._send(command) return command async def image_edits( self, request: Request, image: UploadFile = File(...), # noqa: B008 prompt: str = Form(...), model: str = Form(...), n: int = Form(1), size: str | None = Form(None), response_format: Literal["url", "b64_json"] = Form("b64_json"), input_fidelity: Literal["low", "high"] = Form("low"), stream: str = Form("false"), partial_images: str = Form("0"), quality: Literal["high", "medium", "low"] = Form("medium"), output_format: Literal["png", "jpeg", "webp"] = Form("png"), advanced_params: str | None = Form(None), ) -> ImageGenerationResponse | StreamingResponse: """Handle image editing requests (img2img).""" # Parse string form values to proper types stream_bool = stream.lower() in ("true", "1", "yes") partial_images_int = int(partial_images) if partial_images.isdigit() else 0 parsed_advanced_params: AdvancedImageParams | None = None if advanced_params: with contextlib.suppress(Exception): parsed_advanced_params = AdvancedImageParams.model_validate_json( advanced_params ) command = await self._send_image_edits_command( image=image, prompt=prompt, model=ModelId(model), n=n, size=normalize_image_size(size), response_format=response_format, input_fidelity=input_fidelity, stream=stream_bool, partial_images=partial_images_int, bench=False, quality=quality, output_format=output_format, advanced_params=parsed_advanced_params, ) if stream_bool and partial_images_int > 0: return StreamingResponse( self._generate_image_stream( request=request, command_id=command.command_id, num_images=n, response_format=response_format, ), media_type="text/event-stream", ) return await self._collect_image_generation( request=request, command_id=command.command_id, num_images=n, response_format=response_format, ) async def bench_image_edits( self, request: Request, image: UploadFile = File(...), # noqa: B008 prompt: str = Form(...), model: str = Form(...), n: int = Form(1), size: str | None = Form(None), response_format: Literal["url", "b64_json"] = Form("b64_json"), input_fidelity: Literal["low", "high"] = Form("low"), quality: Literal["high", "medium", "low"] = Form("medium"), output_format: Literal["png", "jpeg", "webp"] = Form("png"), advanced_params: str | None = Form(None), ) -> BenchImageGenerationResponse: """Handle benchmark image editing requests with generation stats.""" parsed_advanced_params: AdvancedImageParams | None = None if advanced_params: with contextlib.suppress(Exception): parsed_advanced_params = AdvancedImageParams.model_validate_json( advanced_params ) command = await self._send_image_edits_command( image=image, prompt=prompt, model=ModelId(model), n=n, size=normalize_image_size(size), response_format=response_format, input_fidelity=input_fidelity, stream=False, partial_images=0, bench=True, quality=quality, output_format=output_format, advanced_params=parsed_advanced_params, ) return await self._collect_image_generation_with_stats( request=request, command_id=command.command_id, num_images=n, response_format=response_format, ) async def claude_messages( self, payload: ClaudeMessagesRequest ) -> ClaudeMessagesResponse | StreamingResponse: """Claude Messages API - adapter.""" task_params = claude_request_to_text_generation(payload) resolved_model = await self._resolve_and_validate_text_model( ModelId(task_params.model) ) task_params = task_params.model_copy(update={"model": resolved_model}) command = TextGeneration(task_params=task_params) await self._send(command) if payload.stream: return StreamingResponse( generate_claude_stream( command.command_id, payload.model, self._token_chunk_stream(command.command_id), ), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "close", "X-Accel-Buffering": "no", }, ) else: return StreamingResponse( collect_claude_response( command.command_id, payload.model, self._token_chunk_stream(command.command_id), ), media_type="application/json", ) async def openai_responses( self, payload: ResponsesRequest ) -> ResponsesResponse | StreamingResponse: """OpenAI Responses API.""" task_params = responses_request_to_text_generation(payload) resolved_model = await self._resolve_and_validate_text_model(task_params.model) task_params = task_params.model_copy(update={"model": resolved_model}) command = TextGeneration(task_params=task_params) await self._send(command) if payload.stream: return StreamingResponse( generate_responses_stream( command.command_id, payload.model, self._token_chunk_stream(command.command_id), ), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "close", "X-Accel-Buffering": "no", }, ) else: return StreamingResponse( collect_responses_response( command.command_id, payload.model, self._token_chunk_stream(command.command_id), ), media_type="application/json", ) async def _ollama_root(self) -> JSONResponse: """Respond to HEAD / from Ollama CLI connectivity checks.""" return JSONResponse(content="Ollama is running") async def ollama_chat( self, request: Request ) -> OllamaChatResponse | StreamingResponse: """Ollama Chat API — accepts JSON regardless of Content-Type.""" body = await request.body() payload = OllamaChatRequest.model_validate_json(body) task_params = ollama_request_to_text_generation(payload) resolved_model = await self._resolve_and_validate_text_model( ModelId(task_params.model) ) task_params = task_params.model_copy(update={"model": resolved_model}) command = TextGeneration(task_params=task_params) await self._send(command) if payload.stream: return StreamingResponse( generate_ollama_chat_stream( command.command_id, self._token_chunk_stream(command.command_id), ), media_type="application/x-ndjson", headers={ "Cache-Control": "no-cache", "Connection": "close", "X-Accel-Buffering": "no", }, ) else: return StreamingResponse( collect_ollama_chat_response( command.command_id, self._token_chunk_stream(command.command_id), ), media_type="application/json", ) async def ollama_generate( self, request: Request ) -> OllamaGenerateResponse | StreamingResponse: """Ollama Generate API — accepts JSON regardless of Content-Type.""" body = await request.body() payload = OllamaGenerateRequest.model_validate_json(body) task_params = ollama_generate_request_to_text_generation(payload) resolved_model = await self._resolve_and_validate_text_model( ModelId(task_params.model) ) task_params = task_params.model_copy(update={"model": resolved_model}) command = TextGeneration(task_params=task_params) await self._send(command) if payload.stream: return StreamingResponse( generate_ollama_generate_stream( command.command_id, self._token_chunk_stream(command.command_id), ), media_type="application/x-ndjson", headers={ "Cache-Control": "no-cache", "Connection": "close", "X-Accel-Buffering": "no", }, ) else: return StreamingResponse( collect_ollama_generate_response( command.command_id, self._token_chunk_stream(command.command_id), ), media_type="application/json", ) async def ollama_tags(self) -> OllamaTagsResponse: """Returns list of models in Ollama tags format. We return the downloaded ones only.""" def none_if_empty(value: str) -> str | None: return value or None downloaded_model_ids: set[str] = set() for node_downloads in self.state.downloads.values(): for dl in node_downloads: if isinstance(dl, DownloadCompleted): downloaded_model_ids.add(dl.shard_metadata.model_card.model_id) cards = [ c for c in await get_model_cards() if c.model_id in downloaded_model_ids ] now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) return OllamaTagsResponse( models=[ OllamaModelTag( name=str(card.model_id), model=str(card.model_id), modified_at=now, size=card.storage_size.in_bytes, digest="sha256:000000000000", details=OllamaModelDetails( family=none_if_empty(card.family), quantization_level=none_if_empty(card.quantization), ), ) for card in cards ] ) async def ollama_show(self, request: Request) -> OllamaShowResponse: """Returns model information in Ollama show format.""" body = await request.body() payload = OllamaShowRequest.model_validate_json(body) model_name = payload.name or payload.model if not model_name: raise HTTPException(status_code=400, detail="name or model is required") try: card = await ModelCard.load(ModelId(model_name)) except Exception as exc: raise HTTPException( status_code=404, detail=f"Model not found: {model_name}" ) from exc return OllamaShowResponse( modelfile=f"FROM {card.model_id}", template="{{ .Prompt }}", details=OllamaModelDetails( family=card.family or None, quantization_level=card.quantization or None, ), ) async def ollama_ps(self) -> OllamaPsResponse: """Returns list of running models (active instances).""" models: list[OllamaPsModel] = [] seen: set[str] = set() for instance in self.state.instances.values(): model_id = str(instance.shard_assignments.model_id) if model_id in seen: continue seen.add(model_id) models.append( OllamaPsModel( name=model_id, model=model_id, size=0, ) ) return OllamaPsResponse(models=models) async def ollama_version(self) -> dict[str, str]: """Returns version information for Ollama API compatibility.""" return {"version": "exo v1.0"} def _calculate_total_available_memory(self) -> Memory: """Calculate total available memory across all nodes in bytes.""" total_available = Memory() for memory in self.state.node_memory.values(): total_available += memory.ram_available return total_available async def get_models(self, status: str | None = Query(default=None)) -> ModelList: """Returns list of available models, optionally filtered by being downloaded.""" cards = await get_model_cards() if status == "downloaded": downloaded_model_ids: set[str] = set() for node_downloads in self.state.downloads.values(): for dl in node_downloads: if isinstance(dl, DownloadCompleted): downloaded_model_ids.add(dl.shard_metadata.model_card.model_id) cards = [c for c in cards if c.model_id in downloaded_model_ids] return ModelList( data=[ ModelListModel( id=card.model_id, hugging_face_id=card.model_id, name=card.model_id.short(), description="", tags=[], storage_size_megabytes=card.storage_size.in_mb, supports_tensor=card.supports_tensor, tasks=[task.value for task in card.tasks], is_custom=is_custom_card(card.model_id), family=card.family, quantization=card.quantization, base_model=card.base_model, capabilities=card.capabilities, ) for card in cards ] ) async def add_custom_model(self, payload: AddCustomModelParams) -> ModelListModel: """Fetch a model from HuggingFace and save as a custom model card.""" try: card = await ModelCard.fetch_from_hf(payload.model_id) except Exception as exc: raise HTTPException( status_code=400, detail=f"Failed to fetch model: {exc}" ) from exc return ModelListModel( id=card.model_id, hugging_face_id=card.model_id, name=card.model_id.short(), description="", tags=[], storage_size_megabytes=int(card.storage_size.in_mb), supports_tensor=card.supports_tensor, tasks=[task.value for task in card.tasks], is_custom=True, ) async def delete_custom_model(self, model_id: ModelId) -> JSONResponse: """Delete a user-added custom model card.""" deleted = await delete_custom_card(model_id) if not deleted: raise HTTPException(status_code=404, detail="Custom model card not found") return JSONResponse( {"message": "Model card deleted", "model_id": str(model_id)} ) async def search_models( self, query: str = "", limit: int = 20 ) -> list[HuggingFaceSearchResult]: """Search HuggingFace Hub — tries mlx-community first, falls back to all of HuggingFace.""" from huggingface_hub import ModelInfo, list_models def _to_results(models: Iterable[ModelInfo]) -> list[HuggingFaceSearchResult]: return [ HuggingFaceSearchResult( id=m.id, author=m.author or "", downloads=m.downloads or 0, likes=m.likes or 0, last_modified=str(m.last_modified or ""), tags=list(m.tags or []), ) for m in models ] # Search mlx-community first mlx_results = _to_results( list_models( search=query or None, author="mlx-community", sort="downloads", limit=limit, ) ) if mlx_results: return mlx_results # Fall back to searching all of HuggingFace return _to_results( list_models( search=query or None, sort="downloads", limit=limit, ) ) async def run(self): shutdown_ev = anyio.Event() try: async with self._tg as tg: logger.info("Starting API") tg.start_soon(self._apply_state) tg.start_soon(self._pause_on_new_election) tg.start_soon(self._cleanup_expired_images) print_startup_banner(self.port) tg.start_soon(self.run_api, shutdown_ev) try: await anyio.sleep_forever() finally: with anyio.CancelScope(shield=True): shutdown_ev.set() finally: self._event_log.close() self.command_sender.close() self.event_receiver.close() async def run_api(self, ev: anyio.Event): cfg = Config() cfg.bind = [f"0.0.0.0:{self.port}"] # nb: shared.logging needs updating if any of this changes cfg.accesslog = None cfg.errorlog = "-" cfg.logger_class = InterceptLogger with anyio.CancelScope(shield=True): await serve( cast(ASGIFramework, self.app), cfg, shutdown_trigger=ev.wait, ) async def _apply_state(self): with self.event_receiver as events: async for i_event in events: self._event_log.append(i_event.event) self.state = apply(self.state, i_event) event = i_event.event if isinstance(event, ChunkGenerated): if queue := self._image_generation_queues.get( event.command_id, None ): assert isinstance(event.chunk, ImageChunk) try: await queue.send(event.chunk) except BrokenResourceError: self._image_generation_queues.pop(event.command_id, None) if queue := self._text_generation_queues.get( event.command_id, None ): assert not isinstance(event.chunk, ImageChunk) try: await queue.send(event.chunk) except BrokenResourceError: self._text_generation_queues.pop(event.command_id, None) if isinstance(event, TracesMerged): self._save_merged_trace(event) def _save_merged_trace(self, event: TracesMerged) -> None: traces = [ TraceEvent( name=t.name, start_us=t.start_us, duration_us=t.duration_us, rank=t.rank, category=t.category, ) for t in event.traces ] output_path = EXO_TRACING_CACHE_DIR / f"trace_{event.task_id}.json" export_trace(traces, output_path) logger.debug(f"Saved merged trace to {output_path}") async def _pause_on_new_election(self): with self.election_receiver as ems: async for message in ems: if message.clock > self.last_completed_election: self.paused = True async def _cleanup_expired_images(self): """Periodically clean up expired images from the store.""" cleanup_interval_seconds = 300 # 5 minutes while True: await anyio.sleep(cleanup_interval_seconds) removed = self._image_store.cleanup_expired() if removed > 0: logger.debug(f"Cleaned up {removed} expired images") async def _send(self, command: Command): while self.paused: await self.paused_ev.wait() await self.command_sender.send( ForwarderCommand(origin=self._system_id, command=command) ) async def _send_download(self, command: DownloadCommand): await self.download_command_sender.send( ForwarderDownloadCommand(origin=self._system_id, command=command) ) async def start_download( self, payload: StartDownloadParams ) -> StartDownloadResponse: command = StartDownload( target_node_id=payload.target_node_id, shard_metadata=payload.shard_metadata, ) await self._send_download(command) return StartDownloadResponse(command_id=command.command_id) async def delete_download( self, node_id: NodeId, model_id: ModelId ) -> DeleteDownloadResponse: command = DeleteDownload( target_node_id=node_id, model_id=ModelId(model_id), ) await self._send_download(command) return DeleteDownloadResponse(command_id=command.command_id) @staticmethod def _get_trace_path(task_id: str) -> Path: trace_path = EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json" if not trace_path.resolve().is_relative_to(EXO_TRACING_CACHE_DIR.resolve()): raise HTTPException(status_code=400, detail=f"Invalid task ID: {task_id}") return trace_path async def list_traces(self) -> TraceListResponse: traces: list[TraceListItem] = [] for trace_file in sorted( EXO_TRACING_CACHE_DIR.glob("trace_*.json"), key=lambda p: p.stat().st_mtime, reverse=True, ): # Extract task_id from filename (trace_{task_id}.json) task_id = trace_file.stem.removeprefix("trace_") stat = trace_file.stat() created_at = datetime.fromtimestamp( stat.st_mtime, tz=timezone.utc ).isoformat() traces.append( TraceListItem( task_id=task_id, created_at=created_at, file_size=stat.st_size, ) ) return TraceListResponse(traces=traces) async def get_trace(self, task_id: str) -> TraceResponse: trace_path = self._get_trace_path(task_id) if not trace_path.exists(): raise HTTPException(status_code=404, detail=f"Trace not found: {task_id}") trace_events = load_trace_file(trace_path) return TraceResponse( task_id=task_id, traces=[ TraceEventResponse( name=event.name, start_us=event.start_us, duration_us=event.duration_us, rank=event.rank, category=event.category, ) for event in trace_events ], ) async def get_trace_stats(self, task_id: str) -> TraceStatsResponse: trace_path = self._get_trace_path(task_id) if not trace_path.exists(): raise HTTPException(status_code=404, detail=f"Trace not found: {task_id}") trace_events = load_trace_file(trace_path) stats = compute_stats(trace_events) return TraceStatsResponse( task_id=task_id, total_wall_time_us=stats.total_wall_time_us, by_category={ category: TraceCategoryStats( total_us=cat_stats.total_us, count=cat_stats.count, min_us=cat_stats.min_us, max_us=cat_stats.max_us, avg_us=cat_stats.avg_us, ) for category, cat_stats in stats.by_category.items() }, by_rank={ rank: TraceRankStats( by_category={ category: TraceCategoryStats( total_us=cat_stats.total_us, count=cat_stats.count, min_us=cat_stats.min_us, max_us=cat_stats.max_us, avg_us=cat_stats.avg_us, ) for category, cat_stats in rank_stats.items() } ) for rank, rank_stats in stats.by_rank.items() }, ) async def get_trace_raw(self, task_id: str) -> FileResponse: trace_path = self._get_trace_path(task_id) if not trace_path.exists(): raise HTTPException(status_code=404, detail=f"Trace not found: {task_id}") return FileResponse( path=trace_path, media_type="application/json", filename=f"trace_{task_id}.json", ) async def delete_traces(self, request: DeleteTracesRequest) -> DeleteTracesResponse: deleted: list[str] = [] not_found: list[str] = [] for task_id in request.task_ids: trace_path = self._get_trace_path(task_id) if trace_path.exists(): trace_path.unlink() deleted.append(task_id) else: not_found.append(task_id) return DeleteTracesResponse(deleted=deleted, not_found=not_found) async def get_onboarding(self) -> JSONResponse: return JSONResponse({"completed": ONBOARDING_COMPLETE_FILE.exists()}) async def complete_onboarding(self) -> JSONResponse: ONBOARDING_COMPLETE_FILE.parent.mkdir(parents=True, exist_ok=True) ONBOARDING_COMPLETE_FILE.write_text("true") return JSONResponse({"completed": True}) ================================================ FILE: src/exo/api/tests/test_api_error_handling.py ================================================ # pyright: reportUnusedFunction=false, reportAny=false from typing import Any from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient from exo.api.main import API def test_http_exception_handler_formats_openai_style() -> None: """Test that HTTPException is converted to OpenAI-style error format.""" app = FastAPI() # Setup exception handler api = object.__new__(API) api.app = app api._setup_exception_handlers() # pyright: ignore[reportPrivateUsage] # Add test routes that raise HTTPException @app.get("/test-error") async def _test_error() -> None: raise HTTPException(status_code=500, detail="Test error message") @app.get("/test-not-found") async def _test_not_found() -> None: raise HTTPException(status_code=404, detail="Resource not found") client = TestClient(app) # Test 500 error response = client.get("/test-error") assert response.status_code == 500 data: dict[str, Any] = response.json() assert "error" in data assert data["error"]["message"] == "Test error message" assert data["error"]["type"] == "Internal Server Error" assert data["error"]["code"] == 500 # Test 404 error response = client.get("/test-not-found") assert response.status_code == 404 data = response.json() assert "error" in data assert data["error"]["message"] == "Resource not found" assert data["error"]["type"] == "Not Found" assert data["error"]["code"] == 404 ================================================ FILE: src/exo/api/tests/test_cancel_command.py ================================================ # pyright: reportUnusedFunction=false, reportAny=false from typing import Any from unittest.mock import AsyncMock, MagicMock from fastapi import FastAPI from fastapi.testclient import TestClient from exo.api.main import API from exo.shared.types.common import CommandId def _make_api() -> Any: """Create a minimal API instance with cancel route and error handler.""" app = FastAPI() api = object.__new__(API) api.app = app api._text_generation_queues = {} # pyright: ignore[reportPrivateUsage] api._image_generation_queues = {} # pyright: ignore[reportPrivateUsage] api._send = AsyncMock() # pyright: ignore[reportPrivateUsage] api._setup_exception_handlers() # pyright: ignore[reportPrivateUsage] app.post("/v1/cancel/{command_id}")(api.cancel_command) return api def test_cancel_nonexistent_command_returns_404() -> None: """Cancel for an unknown command_id returns 404 in OpenAI error format.""" api = _make_api() client = TestClient(api.app) response = client.post("/v1/cancel/nonexistent-id") assert response.status_code == 404 data: dict[str, Any] = response.json() assert "error" in data assert data["error"]["message"] == "Command not found or already completed" assert data["error"]["type"] == "Not Found" assert data["error"]["code"] == 404 def test_cancel_active_text_generation() -> None: """Cancel an active text generation command: returns 200, sender.close() called.""" api = _make_api() client = TestClient(api.app) cid = CommandId("text-cmd-123") sender = MagicMock() api._text_generation_queues[cid] = sender response = client.post(f"/v1/cancel/{cid}") assert response.status_code == 200 data: dict[str, Any] = response.json() assert data["message"] == "Command cancelled." assert data["command_id"] == str(cid) sender.close.assert_called_once() api._send.assert_called_once() task_cancelled = api._send.call_args[0][0] assert task_cancelled.cancelled_command_id == cid def test_cancel_active_image_generation() -> None: """Cancel an active image generation command: returns 200, sender.close() called.""" api = _make_api() client = TestClient(api.app) cid = CommandId("img-cmd-456") sender = MagicMock() api._image_generation_queues[cid] = sender response = client.post(f"/v1/cancel/{cid}") assert response.status_code == 200 data: dict[str, Any] = response.json() assert data["message"] == "Command cancelled." assert data["command_id"] == str(cid) sender.close.assert_called_once() api._send.assert_called_once() task_cancelled = api._send.call_args[0][0] assert task_cancelled.cancelled_command_id == cid ================================================ FILE: src/exo/api/tests/test_claude_api.py ================================================ """Tests for Claude Messages API conversion functions and types.""" import pydantic import pytest from exo.api.adapters.claude import ( claude_request_to_text_generation, finish_reason_to_claude_stop_reason, ) from exo.api.types.claude_api import ( ClaudeMessage, ClaudeMessagesRequest, ClaudeTextBlock, ) from exo.shared.types.common import ModelId class TestFinishReasonToClaudeStopReason: """Tests for finish_reason to Claude stop_reason mapping.""" def test_stop_maps_to_end_turn(self): assert finish_reason_to_claude_stop_reason("stop") == "end_turn" def test_length_maps_to_max_tokens(self): assert finish_reason_to_claude_stop_reason("length") == "max_tokens" def test_tool_calls_maps_to_tool_use(self): assert finish_reason_to_claude_stop_reason("tool_calls") == "tool_use" def test_function_call_maps_to_tool_use(self): assert finish_reason_to_claude_stop_reason("function_call") == "tool_use" def test_content_filter_maps_to_end_turn(self): assert finish_reason_to_claude_stop_reason("content_filter") == "end_turn" def test_none_returns_none(self): assert finish_reason_to_claude_stop_reason(None) is None class TestClaudeRequestToInternal: """Tests for converting Claude Messages API requests to TextGenerationTaskParams.""" def test_basic_request_conversion(self): request = ClaudeMessagesRequest( model=ModelId("claude-3-opus"), max_tokens=100, messages=[ ClaudeMessage(role="user", content="Hello"), ], ) params = claude_request_to_text_generation(request) assert params.model == "claude-3-opus" assert params.max_output_tokens == 100 assert isinstance(params.input, list) assert len(params.input) == 1 assert params.input[0].role == "user" assert params.input[0].content == "Hello" assert params.instructions is None def test_request_with_system_string(self): request = ClaudeMessagesRequest( model=ModelId("claude-3-opus"), max_tokens=100, system="You are a helpful assistant.", messages=[ ClaudeMessage(role="user", content="Hello"), ], ) params = claude_request_to_text_generation(request) assert params.instructions == "You are a helpful assistant." assert isinstance(params.input, list) assert len(params.input) == 1 assert params.input[0].role == "user" assert params.input[0].content == "Hello" def test_request_with_system_text_blocks(self): request = ClaudeMessagesRequest( model=ModelId("claude-3-opus"), max_tokens=100, system=[ ClaudeTextBlock(text="You are helpful. "), ClaudeTextBlock(text="Be concise."), ], messages=[ ClaudeMessage(role="user", content="Hello"), ], ) params = claude_request_to_text_generation(request) assert params.instructions == "You are helpful. Be concise." assert isinstance(params.input, list) assert len(params.input) == 1 def test_request_with_content_blocks(self): request = ClaudeMessagesRequest( model=ModelId("claude-3-opus"), max_tokens=100, messages=[ ClaudeMessage( role="user", content=[ ClaudeTextBlock(text="First part. "), ClaudeTextBlock(text="Second part."), ], ), ], ) params = claude_request_to_text_generation(request) assert isinstance(params.input, list) assert len(params.input) == 1 assert params.input[0].content == "First part. Second part." def test_request_with_multi_turn_conversation(self): request = ClaudeMessagesRequest( model=ModelId("claude-3-opus"), max_tokens=100, messages=[ ClaudeMessage(role="user", content="Hello"), ClaudeMessage(role="assistant", content="Hi there!"), ClaudeMessage(role="user", content="How are you?"), ], ) params = claude_request_to_text_generation(request) assert isinstance(params.input, list) assert len(params.input) == 3 assert params.input[0].role == "user" assert params.input[1].role == "assistant" assert params.input[2].role == "user" def test_request_with_optional_parameters(self): request = ClaudeMessagesRequest( model=ModelId("claude-3-opus"), max_tokens=100, messages=[ClaudeMessage(role="user", content="Hello")], temperature=0.7, top_p=0.9, top_k=40, stop_sequences=["STOP", "END"], stream=True, ) params = claude_request_to_text_generation(request) assert params.temperature == 0.7 assert params.top_p == 0.9 assert params.top_k == 40 assert params.stop == ["STOP", "END"] assert params.stream is True class TestClaudeMessagesRequestValidation: """Tests for Claude Messages API request validation.""" def test_request_requires_model(self): with pytest.raises(pydantic.ValidationError): ClaudeMessagesRequest.model_validate( { "max_tokens": 100, "messages": [{"role": "user", "content": "Hello"}], } ) def test_request_requires_max_tokens(self): with pytest.raises(pydantic.ValidationError): ClaudeMessagesRequest.model_validate( { "model": "claude-3-opus", "messages": [{"role": "user", "content": "Hello"}], } ) def test_request_requires_messages(self): with pytest.raises(pydantic.ValidationError): ClaudeMessagesRequest.model_validate( { "model": "claude-3-opus", "max_tokens": 100, } ) ================================================ FILE: src/exo/api/tests/test_claude_tool_use.py ================================================ """Tests for Claude Messages API tool_use support in the adapter.""" import json from collections.abc import AsyncGenerator from typing import Any, cast from exo.api.adapters.claude import ( ClaudeMessagesResponse, collect_claude_response, generate_claude_stream, ) from exo.api.types import ToolCallItem from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk from exo.shared.types.common import CommandId, ModelId async def _chunks_to_stream( chunks: list[ErrorChunk | ToolCallChunk | TokenChunk], ) -> AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None]: for chunk in chunks: yield chunk async def _collect_response( command_id: CommandId, model: str, chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None], ) -> ClaudeMessagesResponse: """Helper to consume the async generator and parse the JSON response.""" parts: list[str] = [] async for part in collect_claude_response(command_id, model, chunk_stream): parts.append(part) return ClaudeMessagesResponse.model_validate_json("".join(parts)) MODEL = ModelId("test-model") COMMAND_ID = CommandId("cmd_test123") def _parse_sse_events(events: list[str]) -> list[dict[str, Any]]: """Parse SSE event strings into JSON dicts.""" parsed: list[dict[str, Any]] = [] for event_str in events: for line in event_str.strip().split("\n"): if line.startswith("data: "): parsed.append(cast(dict[str, Any], json.loads(line[6:]))) return parsed class TestCollectClaudeResponseToolUse: """Tests for non-streaming tool_use response collection.""" async def test_tool_call_chunk_produces_tool_use_blocks(self): chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [ ToolCallChunk( model=MODEL, usage=None, tool_calls=[ ToolCallItem( name="get_weather", arguments='{"location": "San Francisco"}', ) ], ), ] response = await _collect_response( COMMAND_ID, "test-model", _chunks_to_stream(chunks) ) assert response.stop_reason == "tool_use" tool_blocks = [b for b in response.content if b.type == "tool_use"] assert len(tool_blocks) == 1 block = tool_blocks[0] assert block.type == "tool_use" assert block.name == "get_weather" assert block.input == {"location": "San Francisco"} assert block.id.startswith("toolu_") async def test_multiple_tool_calls(self): chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [ ToolCallChunk( model=MODEL, usage=None, tool_calls=[ ToolCallItem( name="get_weather", arguments='{"location": "SF"}', ), ToolCallItem( name="get_time", arguments='{"timezone": "PST"}', ), ], ), ] response = await _collect_response( COMMAND_ID, "test-model", _chunks_to_stream(chunks) ) assert response.stop_reason == "tool_use" tool_blocks = [b for b in response.content if b.type == "tool_use"] assert len(tool_blocks) == 2 assert tool_blocks[0].name == "get_weather" assert tool_blocks[1].name == "get_time" async def test_mixed_text_and_tool_use(self): chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [ TokenChunk(model=MODEL, text="Let me check ", token_id=1, usage=None), TokenChunk(model=MODEL, text="the weather.", token_id=2, usage=None), ToolCallChunk( model=MODEL, usage=None, tool_calls=[ ToolCallItem( name="get_weather", arguments='{"location": "NYC"}', ) ], ), ] response = await _collect_response( COMMAND_ID, "test-model", _chunks_to_stream(chunks) ) assert response.stop_reason == "tool_use" text_blocks = [b for b in response.content if b.type == "text"] tool_blocks = [b for b in response.content if b.type == "tool_use"] assert len(text_blocks) == 1 assert text_blocks[0].text == "Let me check the weather." assert len(tool_blocks) == 1 assert tool_blocks[0].name == "get_weather" async def test_no_content_produces_empty_text_block(self): chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [] response = await _collect_response( COMMAND_ID, "test-model", _chunks_to_stream(chunks) ) assert len(response.content) == 1 assert response.content[0].type == "text" class TestGenerateClaudeStreamToolUse: """Tests for streaming tool_use event generation.""" async def test_tool_call_emits_tool_use_events(self): chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [ ToolCallChunk( model=MODEL, usage=None, tool_calls=[ ToolCallItem( name="get_weather", arguments='{"location": "SF"}', ) ], ), ] events: list[str] = [] async for event in generate_claude_stream( COMMAND_ID, "test-model", _chunks_to_stream(chunks) ): events.append(event) parsed = _parse_sse_events(events) # Find tool_use content_block_start tool_starts = [ e for e in parsed if e.get("type") == "content_block_start" and cast(dict[str, Any], e.get("content_block", {})).get("type") == "tool_use" ] assert len(tool_starts) == 1 content_block = cast(dict[str, Any], tool_starts[0]["content_block"]) assert content_block["name"] == "get_weather" assert content_block["input"] == {} assert cast(str, content_block["id"]).startswith("toolu_") # Find input_json_delta json_deltas = [ e for e in parsed if e.get("type") == "content_block_delta" and cast(dict[str, Any], e.get("delta", {})).get("type") == "input_json_delta" ] assert len(json_deltas) == 1 delta = cast(dict[str, Any], json_deltas[0]["delta"]) assert json.loads(cast(str, delta["partial_json"])) == {"location": "SF"} # Find message_delta with tool_use stop reason msg_deltas = [e for e in parsed if e.get("type") == "message_delta"] assert len(msg_deltas) == 1 assert cast(dict[str, Any], msg_deltas[0]["delta"])["stop_reason"] == "tool_use" async def test_streaming_mixed_text_and_tool_use(self): chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [ TokenChunk(model=MODEL, text="Hello ", token_id=1, usage=None), ToolCallChunk( model=MODEL, usage=None, tool_calls=[ ToolCallItem( name="search", arguments='{"query": "test"}', ) ], ), ] events: list[str] = [] async for event in generate_claude_stream( COMMAND_ID, "test-model", _chunks_to_stream(chunks) ): events.append(event) parsed = _parse_sse_events(events) # Should have text delta at index 0 text_deltas = [ e for e in parsed if e.get("type") == "content_block_delta" and cast(dict[str, Any], e.get("delta", {})).get("type") == "text_delta" ] assert len(text_deltas) == 1 assert text_deltas[0]["index"] == 0 assert cast(dict[str, Any], text_deltas[0]["delta"])["text"] == "Hello " # Tool block at index 1 tool_starts = [ e for e in parsed if e.get("type") == "content_block_start" and cast(dict[str, Any], e.get("content_block", {})).get("type") == "tool_use" ] assert len(tool_starts) == 1 assert tool_starts[0]["index"] == 1 # Stop reason should be tool_use msg_deltas = [e for e in parsed if e.get("type") == "message_delta"] assert cast(dict[str, Any], msg_deltas[0]["delta"])["stop_reason"] == "tool_use" async def test_streaming_tool_block_stop_events(self): chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = [ ToolCallChunk( model=MODEL, usage=None, tool_calls=[ ToolCallItem(name="fn1", arguments="{}"), ToolCallItem(name="fn2", arguments='{"a": 1}'), ], ), ] events: list[str] = [] async for event in generate_claude_stream( COMMAND_ID, "test-model", _chunks_to_stream(chunks) ): events.append(event) parsed = _parse_sse_events(events) # Two tool block starts (at indices 0 and 1 — no text block when only tools) tool_starts = [ e for e in parsed if e.get("type") == "content_block_start" and cast(dict[str, Any], e.get("content_block", {})).get("type") == "tool_use" ] assert len(tool_starts) == 2 assert tool_starts[0]["index"] == 0 assert tool_starts[1]["index"] == 1 # Two tool block stops (at indices 0 and 1) block_stops = [e for e in parsed if e.get("type") == "content_block_stop"] stop_indices = [e["index"] for e in block_stops] assert 0 in stop_indices assert 1 in stop_indices ================================================ FILE: src/exo/api/tests/test_openai_responses_api.py ================================================ """Tests for OpenAI Responses API wire types. ResponsesRequest is the API wire type for the Responses endpoint. The responses adapter converts it to TextGenerationTaskParams for the pipeline. """ import pydantic import pytest from exo.api.types.openai_responses import ( ResponseInputMessage, ResponsesRequest, ) from exo.shared.types.common import ModelId class TestResponsesRequestValidation: """Tests for OpenAI Responses API request validation.""" def test_request_requires_model(self): with pytest.raises(pydantic.ValidationError): ResponsesRequest.model_validate( { "input": "Hello", } ) def test_request_requires_input(self): with pytest.raises(pydantic.ValidationError): ResponsesRequest.model_validate( { "model": "gpt-4o", } ) def test_request_accepts_string_input(self): request = ResponsesRequest( model=ModelId("gpt-4o"), input="Hello", ) assert request.input == "Hello" def test_request_accepts_message_array_input(self): request = ResponsesRequest( model=ModelId("gpt-4o"), input=[ResponseInputMessage(role="user", content="Hello")], ) assert len(request.input) == 1 ================================================ FILE: src/exo/api/types/__init__.py ================================================ from .api import AddCustomModelParams as AddCustomModelParams from .api import AdvancedImageParams as AdvancedImageParams from .api import BenchChatCompletionRequest as BenchChatCompletionRequest from .api import BenchChatCompletionResponse as BenchChatCompletionResponse from .api import BenchImageGenerationResponse as BenchImageGenerationResponse from .api import BenchImageGenerationTaskParams as BenchImageGenerationTaskParams from .api import CancelCommandResponse as CancelCommandResponse from .api import ChatCompletionChoice as ChatCompletionChoice from .api import ChatCompletionMessage as ChatCompletionMessage from .api import ChatCompletionMessageText as ChatCompletionMessageText from .api import ChatCompletionRequest as ChatCompletionRequest from .api import ChatCompletionResponse as ChatCompletionResponse from .api import CompletionTokensDetails as CompletionTokensDetails from .api import CreateInstanceParams as CreateInstanceParams from .api import CreateInstanceResponse as CreateInstanceResponse from .api import DeleteDownloadResponse as DeleteDownloadResponse from .api import DeleteInstanceResponse as DeleteInstanceResponse from .api import DeleteTracesRequest as DeleteTracesRequest from .api import DeleteTracesResponse as DeleteTracesResponse from .api import ErrorInfo as ErrorInfo from .api import ErrorResponse as ErrorResponse from .api import FinishReason as FinishReason from .api import GenerationStats as GenerationStats from .api import HuggingFaceSearchResult as HuggingFaceSearchResult from .api import ImageData as ImageData from .api import ImageEditsTaskParams as ImageEditsTaskParams from .api import ImageGenerationResponse as ImageGenerationResponse from .api import ImageGenerationStats as ImageGenerationStats from .api import ImageGenerationTaskParams as ImageGenerationTaskParams from .api import ImageListItem as ImageListItem from .api import ImageListResponse as ImageListResponse from .api import ImageSize as ImageSize from .api import Logprobs as Logprobs from .api import LogprobsContentItem as LogprobsContentItem from .api import ModelList as ModelList from .api import ModelListModel as ModelListModel from .api import NodePowerStats as NodePowerStats from .api import PlaceInstanceParams as PlaceInstanceParams from .api import PlacementPreview as PlacementPreview from .api import PlacementPreviewResponse as PlacementPreviewResponse from .api import PowerUsage as PowerUsage from .api import PromptTokensDetails as PromptTokensDetails from .api import StartDownloadParams as StartDownloadParams from .api import StartDownloadResponse as StartDownloadResponse from .api import StreamingChoiceResponse as StreamingChoiceResponse from .api import ToolCall as ToolCall from .api import ToolCallItem as ToolCallItem from .api import TopLogprobItem as TopLogprobItem from .api import TraceCategoryStats as TraceCategoryStats from .api import TraceEventResponse as TraceEventResponse from .api import TraceListItem as TraceListItem from .api import TraceListResponse as TraceListResponse from .api import TraceRankStats as TraceRankStats from .api import TraceResponse as TraceResponse from .api import TraceStatsResponse as TraceStatsResponse from .api import Usage as Usage from .api import normalize_image_size as normalize_image_size ================================================ FILE: src/exo/api/types/api.py ================================================ import time from collections.abc import Generator from typing import Annotated, Any, Literal, get_args from uuid import uuid4 from pydantic import BaseModel, Field, field_validator from exo.shared.models.model_cards import ModelCard, ModelId from exo.shared.types.common import CommandId, NodeId from exo.shared.types.memory import Memory from exo.shared.types.text_generation import ReasoningEffort from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta from exo.shared.types.worker.shards import Sharding, ShardMetadata from exo.utils.pydantic_ext import CamelCaseModel FinishReason = Literal[ "stop", "length", "tool_calls", "content_filter", "function_call", "error" ] class ErrorInfo(BaseModel): message: str type: str param: str | None = None code: int class ErrorResponse(BaseModel): error: ErrorInfo class ModelListModel(BaseModel): id: str object: str = "model" created: int = Field(default_factory=lambda: int(time.time())) owned_by: str = "exo" # openwebui fields hugging_face_id: str = Field(default="") name: str = Field(default="") description: str = Field(default="") context_length: int = Field(default=0) tags: list[str] = Field(default=[]) storage_size_megabytes: int = Field(default=0) supports_tensor: bool = Field(default=False) tasks: list[str] = Field(default=[]) is_custom: bool = Field(default=False) family: str = Field(default="") quantization: str = Field(default="") base_model: str = Field(default="") capabilities: list[str] = Field(default_factory=list) class ModelList(BaseModel): object: Literal["list"] = "list" data: list[ModelListModel] class ChatCompletionMessageText(BaseModel): type: Literal["text"] = "text" text: str class ToolCallItem(BaseModel): id: str = Field(default_factory=lambda: str(uuid4())) name: str arguments: str class ToolCall(BaseModel): id: str index: int | None = None type: Literal["function"] = "function" function: ToolCallItem class ChatCompletionMessage(BaseModel): role: Literal["system", "user", "assistant", "developer", "tool", "function"] content: ( str | ChatCompletionMessageText | list[ChatCompletionMessageText] | None ) = None reasoning_content: str | None = None name: str | None = None tool_calls: list[ToolCall] | None = None tool_call_id: str | None = None function_call: dict[str, Any] | None = None class BenchChatCompletionMessage(ChatCompletionMessage): pass class TopLogprobItem(BaseModel): token: str logprob: float bytes: list[int] | None = None class LogprobsContentItem(BaseModel): token: str logprob: float bytes: list[int] | None = None top_logprobs: list[TopLogprobItem] class Logprobs(BaseModel): content: list[LogprobsContentItem] | None = None class PromptTokensDetails(BaseModel): cached_tokens: int = 0 audio_tokens: int = 0 class CompletionTokensDetails(BaseModel): reasoning_tokens: int = 0 audio_tokens: int = 0 accepted_prediction_tokens: int = 0 rejected_prediction_tokens: int = 0 class Usage(BaseModel): prompt_tokens: int completion_tokens: int total_tokens: int prompt_tokens_details: PromptTokensDetails completion_tokens_details: CompletionTokensDetails class StreamingChoiceResponse(BaseModel): index: int delta: ChatCompletionMessage logprobs: Logprobs | None = None finish_reason: FinishReason | None = None usage: Usage | None = None class ChatCompletionChoice(BaseModel): index: int message: ChatCompletionMessage logprobs: Logprobs | None = None finish_reason: FinishReason | None = None class ChatCompletionResponse(BaseModel): id: str object: Literal["chat.completion"] = "chat.completion" created: int model: str choices: list[ChatCompletionChoice | StreamingChoiceResponse] usage: Usage | None = None service_tier: str | None = None class GenerationStats(BaseModel): prompt_tps: float generation_tps: float prompt_tokens: int generation_tokens: int peak_memory_usage: Memory class ImageGenerationStats(BaseModel): seconds_per_step: float total_generation_time: float num_inference_steps: int num_images: int image_width: int image_height: int peak_memory_usage: Memory class NodePowerStats(BaseModel, frozen=True): node_id: NodeId samples: int avg_sys_power: float class PowerUsage(BaseModel, frozen=True): elapsed_seconds: float nodes: list[NodePowerStats] total_avg_sys_power_watts: float total_energy_joules: float class BenchChatCompletionResponse(ChatCompletionResponse): generation_stats: GenerationStats | None = None power_usage: PowerUsage | None = None class StreamOptions(BaseModel): include_usage: bool = False class ChatCompletionRequest(BaseModel): model: ModelId frequency_penalty: float | None = None messages: list[ChatCompletionMessage] logit_bias: dict[str, int] | None = None logprobs: bool | None = None top_logprobs: int | None = None max_tokens: int | None = None n: int | None = None presence_penalty: float | None = None response_format: dict[str, Any] | None = None seed: int | None = None stop: str | list[str] | None = None stream: bool = False stream_options: StreamOptions | None = None temperature: float | None = None top_p: float | None = None top_k: int | None = None tools: list[dict[str, Any]] | None = None reasoning_effort: ReasoningEffort | None = None enable_thinking: bool | None = None min_p: float | None = None repetition_penalty: float | None = None repetition_context_size: int | None = None tool_choice: str | dict[str, Any] | None = None parallel_tool_calls: bool | None = None user: str | None = None class BenchChatCompletionRequest(ChatCompletionRequest): pass class AddCustomModelParams(BaseModel): model_id: ModelId class HuggingFaceSearchResult(BaseModel): id: str author: str = "" downloads: int = 0 likes: int = 0 last_modified: str = "" tags: list[str] = Field(default_factory=list) class PlaceInstanceParams(BaseModel): model_id: ModelId sharding: Sharding = Sharding.Pipeline instance_meta: InstanceMeta = InstanceMeta.MlxRing min_nodes: int = 1 class CreateInstanceParams(BaseModel): instance: Instance class PlacementPreview(BaseModel): model_id: ModelId sharding: Sharding instance_meta: InstanceMeta instance: Instance | None = None # Keys are NodeId strings, values are additional bytes that would be used on that node memory_delta_by_node: dict[str, int] | None = None error: str | None = None class PlacementPreviewResponse(BaseModel): previews: list[PlacementPreview] class DeleteInstanceTaskParams(BaseModel): instance_id: str class CreateInstanceResponse(BaseModel): message: str command_id: CommandId model_card: ModelCard class DeleteInstanceResponse(BaseModel): message: str command_id: CommandId instance_id: InstanceId class CancelCommandResponse(BaseModel): message: str command_id: CommandId ImageSize = Literal[ "auto", "512x512", "768x768", "1024x768", "768x1024", "1024x1024", "1024x1536", "1536x1024", ] def normalize_image_size(v: object) -> ImageSize: """Shared validator for ImageSize fields: maps None → "auto" and rejects invalid values.""" if v is None: return "auto" if v not in get_args(ImageSize): raise ValueError(f"Invalid size: {v!r}. Must be one of {get_args(ImageSize)}") return v # pyright: ignore[reportReturnType] class AdvancedImageParams(BaseModel): seed: Annotated[int, Field(ge=0)] | None = None num_inference_steps: Annotated[int, Field(ge=1, le=100)] | None = None guidance: Annotated[float, Field(ge=1.0, le=20.0)] | None = None negative_prompt: str | None = None num_sync_steps: Annotated[int, Field(ge=1, le=100)] | None = None class ImageGenerationTaskParams(BaseModel): prompt: str background: str | None = None model: str moderation: str | None = None n: int | None = 1 output_compression: int | None = None output_format: Literal["png", "jpeg", "webp"] = "png" partial_images: int | None = 0 quality: Literal["high", "medium", "low"] | None = "medium" response_format: Literal["url", "b64_json"] | None = "b64_json" size: ImageSize = "auto" stream: bool | None = False style: str | None = "vivid" user: str | None = None advanced_params: AdvancedImageParams | None = None # Internal flag for benchmark mode - set by API, preserved through serialization bench: bool = False @field_validator("size", mode="before") @classmethod def normalize_size(cls, v: object) -> ImageSize: return normalize_image_size(v) class BenchImageGenerationTaskParams(ImageGenerationTaskParams): bench: bool = True class ImageEditsTaskParams(BaseModel): """Internal task params for image-editing requests.""" image_data: str = "" # Base64-encoded image (empty when using chunked transfer) total_input_chunks: int = 0 prompt: str model: str n: int | None = 1 quality: Literal["high", "medium", "low"] | None = "medium" output_format: Literal["png", "jpeg", "webp"] = "png" response_format: Literal["url", "b64_json"] | None = "b64_json" size: ImageSize = "auto" image_strength: float | None = 0.7 stream: bool = False partial_images: int | None = 0 advanced_params: AdvancedImageParams | None = None bench: bool = False @field_validator("size", mode="before") @classmethod def normalize_size(cls, v: object) -> ImageSize: return normalize_image_size(v) def __repr_args__(self) -> Generator[tuple[str, Any], None, None]: for name, value in super().__repr_args__(): # pyright: ignore[reportAny] if name == "image_data": yield name, f"<{len(self.image_data)} chars>" elif name is not None: yield name, value class ImageData(BaseModel): b64_json: str | None = None url: str | None = None revised_prompt: str | None = None def __repr_args__(self) -> Generator[tuple[str, Any], None, None]: for name, value in super().__repr_args__(): # pyright: ignore[reportAny] if name == "b64_json" and self.b64_json is not None: yield name, f"<{len(self.b64_json)} chars>" elif name is not None: yield name, value class ImageGenerationResponse(BaseModel): created: int = Field(default_factory=lambda: int(time.time())) data: list[ImageData] class BenchImageGenerationResponse(ImageGenerationResponse): generation_stats: ImageGenerationStats | None = None power_usage: PowerUsage | None = None class ImageListItem(BaseModel, frozen=True): image_id: str url: str content_type: str expires_at: float class ImageListResponse(BaseModel, frozen=True): data: list[ImageListItem] class StartDownloadParams(CamelCaseModel): target_node_id: NodeId shard_metadata: ShardMetadata class StartDownloadResponse(CamelCaseModel): command_id: CommandId class DeleteDownloadResponse(CamelCaseModel): command_id: CommandId class TraceEventResponse(CamelCaseModel): name: str start_us: int duration_us: int rank: int category: str class TraceResponse(CamelCaseModel): task_id: str traces: list[TraceEventResponse] class TraceCategoryStats(CamelCaseModel): total_us: int count: int min_us: int max_us: int avg_us: float class TraceRankStats(CamelCaseModel): by_category: dict[str, TraceCategoryStats] class TraceStatsResponse(CamelCaseModel): task_id: str total_wall_time_us: int by_category: dict[str, TraceCategoryStats] by_rank: dict[int, TraceRankStats] class TraceListItem(CamelCaseModel): task_id: str created_at: str file_size: int class TraceListResponse(CamelCaseModel): traces: list[TraceListItem] class DeleteTracesRequest(CamelCaseModel): task_ids: list[str] class DeleteTracesResponse(CamelCaseModel): deleted: list[str] not_found: list[str] ================================================ FILE: src/exo/api/types/claude_api.py ================================================ """Claude Messages API types for request/response conversion.""" from typing import Any, Literal from pydantic import BaseModel, Field from exo.shared.types.common import ModelId # Tool definition types ClaudeToolInputSchema = dict[str, Any] class ClaudeToolDefinition(BaseModel, frozen=True): """Tool definition in Claude Messages API request.""" name: str description: str | None = None input_schema: ClaudeToolInputSchema # Type aliases ClaudeRole = Literal["user", "assistant"] ClaudeStopReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] # Content block types class ClaudeTextBlock(BaseModel, frozen=True): """Text content block in Claude Messages API.""" type: Literal["text"] = "text" text: str class ClaudeImageSource(BaseModel, frozen=True): """Image source for Claude image blocks.""" type: Literal["base64", "url"] media_type: str | None = None data: str | None = None url: str | None = None class ClaudeImageBlock(BaseModel, frozen=True): """Image content block in Claude Messages API.""" type: Literal["image"] = "image" source: ClaudeImageSource class ClaudeThinkingBlock(BaseModel, frozen=True): """Thinking content block in Claude Messages API.""" type: Literal["thinking"] = "thinking" thinking: str signature: str | None = None class ClaudeToolUseBlock(BaseModel, frozen=True): """Tool use content block in Claude Messages API.""" type: Literal["tool_use"] = "tool_use" id: str name: str input: dict[str, Any] class ClaudeToolResultBlock(BaseModel, frozen=True): """Tool result content block in Claude Messages API request.""" type: Literal["tool_result"] = "tool_result" tool_use_id: str content: str | list[ClaudeTextBlock] | None = None is_error: bool | None = None cache_control: dict[str, str] | None = None ClaudeContentBlock = ( ClaudeTextBlock | ClaudeImageBlock | ClaudeThinkingBlock | ClaudeToolUseBlock ) # Input content blocks can also include tool_result (sent by user after tool_use) ClaudeInputContentBlock = ( ClaudeTextBlock | ClaudeImageBlock | ClaudeThinkingBlock | ClaudeToolUseBlock | ClaudeToolResultBlock ) # Request types class ClaudeMessage(BaseModel, frozen=True): """Message in Claude Messages API request.""" role: ClaudeRole content: str | list[ClaudeInputContentBlock] class ClaudeThinkingConfig(BaseModel, frozen=True): type: Literal["enabled", "disabled", "adaptive"] budget_tokens: int | None = None class ClaudeMessagesRequest(BaseModel): """Request body for Claude Messages API.""" model: ModelId max_tokens: int messages: list[ClaudeMessage] system: str | list[ClaudeTextBlock] | None = None stop_sequences: list[str] | None = None stream: bool = False temperature: float | None = None top_p: float | None = None top_k: int | None = None tools: list[ClaudeToolDefinition] | None = None metadata: dict[str, str] | None = None thinking: ClaudeThinkingConfig | None = None # Response types class ClaudeUsage(BaseModel, frozen=True): """Token usage in Claude Messages API response.""" input_tokens: int output_tokens: int class ClaudeMessagesResponse(BaseModel, frozen=True): """Response body for Claude Messages API.""" id: str type: Literal["message"] = "message" role: Literal["assistant"] = "assistant" content: list[ClaudeContentBlock] model: str stop_reason: ClaudeStopReason | None = None stop_sequence: str | None = None usage: ClaudeUsage # Streaming event types class ClaudeMessageStart(BaseModel, frozen=True): """Partial message in message_start event.""" id: str type: Literal["message"] = "message" role: Literal["assistant"] = "assistant" content: list[ClaudeTextBlock] = Field(default_factory=list) model: str stop_reason: ClaudeStopReason | None = None stop_sequence: str | None = None usage: ClaudeUsage class ClaudeMessageStartEvent(BaseModel, frozen=True): """Event sent at start of message stream.""" type: Literal["message_start"] = "message_start" message: ClaudeMessageStart class ClaudeContentBlockStartEvent(BaseModel, frozen=True): """Event sent at start of a content block.""" type: Literal["content_block_start"] = "content_block_start" index: int content_block: ClaudeTextBlock | ClaudeThinkingBlock | ClaudeToolUseBlock class ClaudeTextDelta(BaseModel, frozen=True): """Delta for text content block.""" type: Literal["text_delta"] = "text_delta" text: str class ClaudeThinkingDelta(BaseModel, frozen=True): """Delta for thinking content block.""" type: Literal["thinking_delta"] = "thinking_delta" thinking: str class ClaudeInputJsonDelta(BaseModel, frozen=True): """Delta for tool use input JSON content block.""" type: Literal["input_json_delta"] = "input_json_delta" partial_json: str class ClaudeContentBlockDeltaEvent(BaseModel, frozen=True): """Event sent for content block delta.""" type: Literal["content_block_delta"] = "content_block_delta" index: int delta: ClaudeTextDelta | ClaudeThinkingDelta | ClaudeInputJsonDelta class ClaudeContentBlockStopEvent(BaseModel, frozen=True): """Event sent at end of a content block.""" type: Literal["content_block_stop"] = "content_block_stop" index: int class ClaudeMessageDeltaUsage(BaseModel, frozen=True): """Usage in message_delta event.""" output_tokens: int class ClaudeMessageDelta(BaseModel, frozen=True): """Delta in message_delta event.""" stop_reason: ClaudeStopReason | None = None stop_sequence: str | None = None class ClaudeMessageDeltaEvent(BaseModel, frozen=True): """Event sent with final message delta.""" type: Literal["message_delta"] = "message_delta" delta: ClaudeMessageDelta usage: ClaudeMessageDeltaUsage class ClaudeMessageStopEvent(BaseModel, frozen=True): """Event sent at end of message stream.""" type: Literal["message_stop"] = "message_stop" ClaudeStreamEvent = ( ClaudeMessageStartEvent | ClaudeContentBlockStartEvent | ClaudeContentBlockDeltaEvent | ClaudeContentBlockStopEvent | ClaudeMessageDeltaEvent | ClaudeMessageStopEvent ) ================================================ FILE: src/exo/api/types/ollama_api.py ================================================ from __future__ import annotations import time from typing import Any, Literal from pydantic import BaseModel, Field from exo.shared.models.model_cards import ModelId # https://github.com/ollama/ollama/blob/main/docs/api.md OllamaRole = Literal["system", "user", "assistant", "tool"] OllamaDoneReason = Literal["stop", "length", "tool_call", "error"] class OllamaToolFunction(BaseModel, frozen=True): name: str arguments: dict[str, Any] | str index: int | None = None class OllamaToolCall(BaseModel, frozen=True): id: str | None = None type: Literal["function"] | None = None function: OllamaToolFunction class OllamaMessage(BaseModel, frozen=True): role: OllamaRole content: str | None = None thinking: str | None = None tool_calls: list[OllamaToolCall] | None = None name: str | None = None tool_name: str | None = None images: list[str] | None = None class OllamaOptions(BaseModel, frozen=True): num_predict: int | None = None temperature: float | None = None top_p: float | None = None top_k: int | None = None stop: str | list[str] | None = None seed: int | None = None class OllamaChatRequest(BaseModel, frozen=True): model: ModelId messages: list[OllamaMessage] stream: bool = True options: OllamaOptions | None = None tools: list[dict[str, Any]] | None = None format: Literal["json"] | dict[str, Any] | None = None keep_alive: str | int | None = None think: bool | None = None class OllamaGenerateRequest(BaseModel, frozen=True): model: ModelId prompt: str = "" system: str | None = None stream: bool = True options: OllamaOptions | None = None format: Literal["json"] | dict[str, Any] | None = None keep_alive: str | int | None = None think: bool | None = None raw: bool = False class OllamaGenerateResponse(BaseModel, frozen=True, strict=True): model: str created_at: str = Field( default_factory=lambda: time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) ) response: str thinking: str | None = None done: bool done_reason: OllamaDoneReason | None = None total_duration: int | None = None load_duration: int | None = None prompt_eval_count: int | None = None prompt_eval_duration: int | None = None eval_count: int | None = None eval_duration: int | None = None class OllamaShowRequest(BaseModel, frozen=True): name: str | None = None model: str | None = None verbose: bool | None = None class OllamaChatResponse(BaseModel, frozen=True, strict=True): model: str created_at: str = Field( default_factory=lambda: time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) ) message: OllamaMessage done: bool done_reason: OllamaDoneReason | None = None total_duration: int | None = None load_duration: int | None = None prompt_eval_count: int | None = None prompt_eval_duration: int | None = None eval_count: int | None = None eval_duration: int | None = None class OllamaModelDetails(BaseModel, frozen=True, strict=True): format: str | None = None family: str | None = None parameter_size: str | None = None quantization_level: str | None = None class OllamaModelTag(BaseModel, frozen=True, strict=True): name: str model: str | None = None modified_at: str | None = None size: int | None = None digest: str | None = None details: OllamaModelDetails | None = None class OllamaTagsResponse(BaseModel, frozen=True, strict=True): models: list[OllamaModelTag] class OllamaShowResponse(BaseModel, frozen=True, strict=True): modelfile: str | None = None parameters: str | None = None template: str | None = None details: OllamaModelDetails | None = None model_info: dict[str, Any] | None = None class OllamaPsModel(BaseModel, frozen=True, strict=True): name: str model: str size: int digest: str | None = None details: OllamaModelDetails | None = None expires_at: str | None = None size_vram: int | None = None class OllamaPsResponse(BaseModel, frozen=True, strict=True): models: list[OllamaPsModel] ================================================ FILE: src/exo/api/types/openai_responses.py ================================================ """OpenAI Responses API wire types. These types model the OpenAI Responses API request/response format. ResponsesRequest is the API-level wire type; for the canonical internal task params type used by the inference pipeline, see ``exo.shared.types.text_generation.TextGenerationTaskParams``. """ import time from typing import Any, Literal from pydantic import BaseModel, Field from exo.shared.types.common import ModelId from exo.shared.types.text_generation import ReasoningEffort # Type aliases ResponseStatus = Literal["completed", "failed", "in_progress", "incomplete"] ResponseRole = Literal["user", "assistant", "system", "developer"] # Request input content part types class ResponseInputTextPart(BaseModel, frozen=True): """Text content part in a Responses API input message.""" type: Literal["input_text"] = "input_text" text: str class ResponseOutputTextPart(BaseModel, frozen=True): """Output text content part (used when replaying assistant messages in input).""" type: Literal["output_text"] = "output_text" text: str ResponseContentPart = ResponseInputTextPart | ResponseOutputTextPart # Request input item types class ResponseInputMessage(BaseModel, frozen=True): """Input message for Responses API.""" role: ResponseRole content: str | list[ResponseContentPart] type: Literal["message"] = "message" class FunctionCallInputItem(BaseModel, frozen=True): """Function call item replayed in input (from a previous assistant response).""" type: Literal["function_call"] = "function_call" id: str | None = None call_id: str name: str arguments: str status: ResponseStatus | None = None class FunctionCallOutputInputItem(BaseModel, frozen=True): """Function call output item in input (user providing tool results).""" type: Literal["function_call_output"] = "function_call_output" call_id: str output: str id: str | None = None status: ResponseStatus | None = None ResponseInputItem = ( ResponseInputMessage | FunctionCallInputItem | FunctionCallOutputInputItem ) class Reasoning(BaseModel, frozen=True): """Reasoning configuration for OpenAI Responses API.""" effort: ReasoningEffort | None = None summary: Literal["auto", "concise", "detailed"] | None = None class ResponsesRequest(BaseModel, frozen=True): """Request body for OpenAI Responses API. This is the API wire type for the Responses endpoint. The canonical internal task params type is ``TextGenerationTaskParams``; see the ``responses_request_to_text_generation`` adapter for conversion. """ # --- OpenAI Responses API standard fields --- model: ModelId input: str | list[ResponseInputItem] instructions: str | None = None max_output_tokens: int | None = None temperature: float | None = None top_p: float | None = None stream: bool = False tools: list[dict[str, Any]] | None = None metadata: dict[str, str] | None = None reasoning: Reasoning | None = None # --- exo extensions (not in OpenAI Responses API spec) --- enable_thinking: bool | None = Field( default=None, description="[exo extension] Boolean thinking toggle. Not part of the OpenAI Responses API.", json_schema_extra={"x-exo-extension": True}, ) top_k: int | None = Field( default=None, description="[exo extension] Top-k sampling parameter. Not part of the OpenAI Responses API.", json_schema_extra={"x-exo-extension": True}, ) stop: str | list[str] | None = Field( default=None, description="[exo extension] Stop sequence(s). Not part of the OpenAI Responses API.", json_schema_extra={"x-exo-extension": True}, ) seed: int | None = Field( default=None, description="[exo extension] Seed for deterministic sampling. Not part of the OpenAI Responses API.", json_schema_extra={"x-exo-extension": True}, ) # --- Internal fields (preserved during serialization, hidden from OpenAPI schema) --- chat_template_messages: list[dict[str, Any]] | None = Field( default=None, description="Internal: pre-formatted messages for tokenizer chat template. Not part of the OpenAI Responses API.", json_schema_extra={"x-exo-internal": True}, ) # Response types class ResponseOutputText(BaseModel, frozen=True): """Text content in response output.""" type: Literal["output_text"] = "output_text" text: str annotations: list[dict[str, str]] = Field(default_factory=list) class ResponseMessageItem(BaseModel, frozen=True): """Message item in response output array.""" type: Literal["message"] = "message" id: str role: Literal["assistant"] = "assistant" content: list[ResponseOutputText] status: ResponseStatus = "completed" class ResponseFunctionCallItem(BaseModel, frozen=True): """Function call item in response output array.""" type: Literal["function_call"] = "function_call" id: str call_id: str name: str arguments: str status: ResponseStatus = "completed" class ResponseReasoningSummaryText(BaseModel, frozen=True): """Summary text part in a reasoning output item.""" type: Literal["summary_text"] = "summary_text" text: str class ResponseReasoningItem(BaseModel, frozen=True): """Reasoning output item in response output array.""" type: Literal["reasoning"] = "reasoning" id: str summary: list[ResponseReasoningSummaryText] = Field(default_factory=list) status: ResponseStatus = "completed" ResponseItem = ResponseMessageItem | ResponseFunctionCallItem | ResponseReasoningItem class ResponseUsage(BaseModel, frozen=True): """Token usage in Responses API response.""" input_tokens: int output_tokens: int total_tokens: int class ResponsesResponse(BaseModel, frozen=True): """Response body for OpenAI Responses API.""" id: str object: Literal["response"] = "response" created_at: int = Field(default_factory=lambda: int(time.time())) status: ResponseStatus = "completed" model: str output: list[ResponseItem] output_text: str usage: ResponseUsage | None = None # Streaming event types class ResponseCreatedEvent(BaseModel, frozen=True): """Event sent when response is created.""" type: Literal["response.created"] = "response.created" sequence_number: int response: ResponsesResponse class ResponseInProgressEvent(BaseModel, frozen=True): """Event sent when response starts processing.""" type: Literal["response.in_progress"] = "response.in_progress" sequence_number: int response: ResponsesResponse class ResponseOutputItemAddedEvent(BaseModel, frozen=True): """Event sent when an output item is added.""" type: Literal["response.output_item.added"] = "response.output_item.added" sequence_number: int output_index: int item: ResponseItem class ResponseContentPartAddedEvent(BaseModel, frozen=True): """Event sent when a content part is added.""" type: Literal["response.content_part.added"] = "response.content_part.added" sequence_number: int item_id: str output_index: int content_index: int part: ResponseOutputText class ResponseTextDeltaEvent(BaseModel, frozen=True): """Event sent for text delta during streaming.""" type: Literal["response.output_text.delta"] = "response.output_text.delta" sequence_number: int item_id: str output_index: int content_index: int delta: str class ResponseTextDoneEvent(BaseModel, frozen=True): """Event sent when text content is done.""" type: Literal["response.output_text.done"] = "response.output_text.done" sequence_number: int item_id: str output_index: int content_index: int text: str class ResponseContentPartDoneEvent(BaseModel, frozen=True): """Event sent when a content part is done.""" type: Literal["response.content_part.done"] = "response.content_part.done" sequence_number: int item_id: str output_index: int content_index: int part: ResponseOutputText class ResponseOutputItemDoneEvent(BaseModel, frozen=True): """Event sent when an output item is done.""" type: Literal["response.output_item.done"] = "response.output_item.done" sequence_number: int output_index: int item: ResponseItem class ResponseFunctionCallArgumentsDeltaEvent(BaseModel, frozen=True): """Event sent for function call arguments delta during streaming.""" type: Literal["response.function_call_arguments.delta"] = ( "response.function_call_arguments.delta" ) sequence_number: int item_id: str output_index: int delta: str class ResponseFunctionCallArgumentsDoneEvent(BaseModel, frozen=True): """Event sent when function call arguments are complete.""" type: Literal["response.function_call_arguments.done"] = ( "response.function_call_arguments.done" ) sequence_number: int item_id: str output_index: int name: str arguments: str class ResponseReasoningSummaryPartAddedEvent(BaseModel, frozen=True): """Event sent when a reasoning summary part is added.""" type: Literal["response.reasoning_summary_part.added"] = ( "response.reasoning_summary_part.added" ) sequence_number: int item_id: str output_index: int summary_index: int part: ResponseReasoningSummaryText class ResponseReasoningSummaryTextDeltaEvent(BaseModel, frozen=True): """Event sent for reasoning summary text delta during streaming.""" type: Literal["response.reasoning_summary_text.delta"] = ( "response.reasoning_summary_text.delta" ) sequence_number: int item_id: str output_index: int summary_index: int delta: str class ResponseReasoningSummaryTextDoneEvent(BaseModel, frozen=True): """Event sent when reasoning summary text is done.""" type: Literal["response.reasoning_summary_text.done"] = ( "response.reasoning_summary_text.done" ) sequence_number: int item_id: str output_index: int summary_index: int text: str class ResponseReasoningSummaryPartDoneEvent(BaseModel, frozen=True): """Event sent when a reasoning summary part is done.""" type: Literal["response.reasoning_summary_part.done"] = ( "response.reasoning_summary_part.done" ) sequence_number: int item_id: str output_index: int summary_index: int part: ResponseReasoningSummaryText class ResponseCompletedEvent(BaseModel, frozen=True): """Event sent when response is completed.""" type: Literal["response.completed"] = "response.completed" sequence_number: int response: ResponsesResponse ResponsesStreamEvent = ( ResponseCreatedEvent | ResponseInProgressEvent | ResponseOutputItemAddedEvent | ResponseContentPartAddedEvent | ResponseTextDeltaEvent | ResponseTextDoneEvent | ResponseContentPartDoneEvent | ResponseOutputItemDoneEvent | ResponseFunctionCallArgumentsDeltaEvent | ResponseFunctionCallArgumentsDoneEvent | ResponseReasoningSummaryPartAddedEvent | ResponseReasoningSummaryTextDeltaEvent | ResponseReasoningSummaryTextDoneEvent | ResponseReasoningSummaryPartDoneEvent | ResponseCompletedEvent ) ================================================ FILE: src/exo/download/coordinator.py ================================================ from dataclasses import dataclass, field import anyio from anyio import current_time from loguru import logger from exo.download.download_utils import ( RepoDownloadProgress, delete_model, map_repo_download_progress_to_download_progress_data, resolve_model_in_path, ) from exo.download.shard_downloader import ShardDownloader from exo.shared.constants import EXO_MODELS_DIR, EXO_MODELS_PATH from exo.shared.models.model_cards import ModelId, get_model_cards from exo.shared.types.commands import ( CancelDownload, DeleteDownload, ForwarderDownloadCommand, StartDownload, ) from exo.shared.types.common import NodeId from exo.shared.types.events import ( Event, NodeDownloadProgress, ) from exo.shared.types.worker.downloads import ( DownloadCompleted, DownloadFailed, DownloadOngoing, DownloadPending, DownloadProgress, ) from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata from exo.utils.channels import Receiver, Sender from exo.utils.task_group import TaskGroup @dataclass class DownloadCoordinator: node_id: NodeId shard_downloader: ShardDownloader download_command_receiver: Receiver[ForwarderDownloadCommand] event_sender: Sender[Event] offline: bool = False # Local state download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict) active_downloads: dict[ModelId, anyio.CancelScope] = field(default_factory=dict) _tg: TaskGroup = field(init=False, default_factory=TaskGroup) # Per-model throttle for download progress events _last_progress_time: dict[ModelId, float] = field(default_factory=dict) def __post_init__(self) -> None: self.shard_downloader.on_progress(self._download_progress_callback) def _model_dir(self, model_id: ModelId) -> str: return str(EXO_MODELS_DIR / model_id.normalize()) async def _download_progress_callback( self, callback_shard: ShardMetadata, progress: RepoDownloadProgress ) -> None: model_id = callback_shard.model_card.model_id throttle_interval_secs = 1.0 if progress.status == "complete": completed = DownloadCompleted( shard_metadata=callback_shard, node_id=self.node_id, total=progress.total, model_directory=self._model_dir(model_id), ) self.download_status[model_id] = completed await self.event_sender.send( NodeDownloadProgress(download_progress=completed) ) self._last_progress_time.pop(model_id, None) elif ( progress.status == "in_progress" and current_time() - self._last_progress_time.get(model_id, 0.0) > throttle_interval_secs ): ongoing = DownloadOngoing( node_id=self.node_id, shard_metadata=callback_shard, download_progress=map_repo_download_progress_to_download_progress_data( progress ), model_directory=self._model_dir(model_id), ) self.download_status[model_id] = ongoing await self.event_sender.send( NodeDownloadProgress(download_progress=ongoing) ) self._last_progress_time[model_id] = current_time() async def run(self) -> None: logger.info( f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}" ) async with self._tg as tg: tg.start_soon(self._command_processor) tg.start_soon(self._emit_existing_download_progress) def shutdown(self) -> None: self._tg.cancel_tasks() async def _command_processor(self) -> None: with self.download_command_receiver as commands: async for cmd in commands: # Only process commands targeting this node if cmd.command.target_node_id != self.node_id: continue match cmd.command: case StartDownload(shard_metadata=shard): await self._start_download(shard) case DeleteDownload(model_id=model_id): await self._delete_download(model_id) case CancelDownload(model_id=model_id): await self._cancel_download(model_id) async def _cancel_download(self, model_id: ModelId) -> None: if model_id in self.active_downloads and model_id in self.download_status: logger.info(f"Cancelling download for {model_id}") self.active_downloads[model_id].cancel() current_status = self.download_status[model_id] pending = DownloadPending( shard_metadata=current_status.shard_metadata, node_id=self.node_id, model_directory=self._model_dir(model_id), ) self.download_status[model_id] = pending await self.event_sender.send( NodeDownloadProgress(download_progress=pending) ) async def _start_download(self, shard: ShardMetadata) -> None: model_id = shard.model_card.model_id # Check if already downloading, complete, or recently failed if model_id in self.download_status: status = self.download_status[model_id] if isinstance(status, (DownloadOngoing, DownloadCompleted, DownloadFailed)): logger.debug( f"Download for {model_id} already in progress, complete, or failed, skipping" ) return # Check EXO_MODELS_PATH for pre-downloaded models found_path = resolve_model_in_path(model_id) if found_path is not None: logger.info( f"DownloadCoordinator: Model {model_id} found in EXO_MODELS_PATH at {found_path}" ) completed = DownloadCompleted( shard_metadata=shard, node_id=self.node_id, total=shard.model_card.storage_size, model_directory=str(found_path), read_only=True, ) self.download_status[model_id] = completed await self.event_sender.send( NodeDownloadProgress(download_progress=completed) ) return # Emit pending status progress = DownloadPending( shard_metadata=shard, node_id=self.node_id, model_directory=self._model_dir(model_id), ) self.download_status[model_id] = progress await self.event_sender.send(NodeDownloadProgress(download_progress=progress)) # Check initial status from downloader initial_progress = ( await self.shard_downloader.get_shard_download_status_for_shard(shard) ) if initial_progress.status == "complete": completed = DownloadCompleted( shard_metadata=shard, node_id=self.node_id, total=initial_progress.total, model_directory=self._model_dir(model_id), ) self.download_status[model_id] = completed await self.event_sender.send( NodeDownloadProgress(download_progress=completed) ) return if self.offline: logger.warning( f"Offline mode: model {model_id} is not fully available locally, cannot download" ) failed = DownloadFailed( shard_metadata=shard, node_id=self.node_id, error_message=f"Model files not found locally in offline mode: {model_id}", model_directory=self._model_dir(model_id), ) self.download_status[model_id] = failed await self.event_sender.send(NodeDownloadProgress(download_progress=failed)) return # Start actual download self._start_download_task(shard, initial_progress) def _start_download_task( self, shard: ShardMetadata, initial_progress: RepoDownloadProgress ) -> None: model_id = shard.model_card.model_id # Emit ongoing status status = DownloadOngoing( node_id=self.node_id, shard_metadata=shard, download_progress=map_repo_download_progress_to_download_progress_data( initial_progress ), model_directory=self._model_dir(model_id), ) self.download_status[model_id] = status self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status)) async def download_wrapper(cancel_scope: anyio.CancelScope) -> None: try: with cancel_scope: await self.shard_downloader.ensure_shard(shard) except Exception as e: logger.error(f"Download failed for {model_id}: {e}") failed = DownloadFailed( shard_metadata=shard, node_id=self.node_id, error_message=str(e), model_directory=self._model_dir(model_id), ) self.download_status[model_id] = failed await self.event_sender.send( NodeDownloadProgress(download_progress=failed) ) except anyio.get_cancelled_exc_class(): # ignore cancellation - let cleanup do its thing pass finally: self.active_downloads.pop(model_id, None) scope = anyio.CancelScope() self._tg.start_soon(download_wrapper, scope) self.active_downloads[model_id] = scope async def _delete_download(self, model_id: ModelId) -> None: # Protect read-only models (from EXO_MODELS_PATH) from deletion if model_id in self.download_status: current = self.download_status[model_id] if isinstance(current, DownloadCompleted) and current.read_only: logger.warning( f"Refusing to delete read-only model {model_id} (from EXO_MODELS_PATH)" ) return # Cancel if active if model_id in self.active_downloads: logger.info(f"Cancelling active download for {model_id} before deletion") self.active_downloads[model_id].cancel() # Delete from disk logger.info(f"Deleting model files for {model_id}") deleted = await delete_model(model_id) if deleted: logger.info(f"Successfully deleted model {model_id}") else: logger.warning(f"Model {model_id} was not found on disk") # Emit pending status to reset UI state, then remove from local tracking if model_id in self.download_status: current_status = self.download_status[model_id] pending = DownloadPending( shard_metadata=current_status.shard_metadata, node_id=self.node_id, model_directory=self._model_dir(model_id), ) await self.event_sender.send( NodeDownloadProgress(download_progress=pending) ) del self.download_status[model_id] async def _emit_existing_download_progress(self) -> None: while True: try: logger.debug( "DownloadCoordinator: Fetching and emitting existing download progress..." ) async for ( _, progress, ) in self.shard_downloader.get_shard_download_status(): model_id = progress.shard.model_card.model_id # Active downloads emit progress via the callback — don't overwrite if model_id in self.active_downloads: continue if progress.status == "complete": status: DownloadProgress = DownloadCompleted( node_id=self.node_id, shard_metadata=progress.shard, total=progress.total, model_directory=self._model_dir( progress.shard.model_card.model_id ), ) elif progress.status in ["in_progress", "not_started"]: if progress.downloaded_this_session.in_bytes == 0: status = DownloadPending( node_id=self.node_id, shard_metadata=progress.shard, model_directory=self._model_dir( progress.shard.model_card.model_id ), downloaded=progress.downloaded, total=progress.total, ) else: status = DownloadOngoing( node_id=self.node_id, shard_metadata=progress.shard, download_progress=map_repo_download_progress_to_download_progress_data( progress ), model_directory=self._model_dir( progress.shard.model_card.model_id ), ) else: continue self.download_status[progress.shard.model_card.model_id] = status await self.event_sender.send( NodeDownloadProgress(download_progress=status) ) # Scan EXO_MODELS_PATH for pre-downloaded models if EXO_MODELS_PATH is not None: for card in await get_model_cards(): mid = card.model_id if mid in self.active_downloads: continue if isinstance( self.download_status.get(mid), (DownloadCompleted, DownloadOngoing, DownloadFailed), ): continue found = resolve_model_in_path(mid) if found is not None: path_shard = PipelineShardMetadata( model_card=card, device_rank=0, world_size=1, start_layer=0, end_layer=card.n_layers, n_layers=card.n_layers, ) path_completed: DownloadProgress = DownloadCompleted( node_id=self.node_id, shard_metadata=path_shard, total=card.storage_size, model_directory=str(found), read_only=True, ) self.download_status[mid] = path_completed await self.event_sender.send( NodeDownloadProgress(download_progress=path_completed) ) logger.debug( "DownloadCoordinator: Done emitting existing download progress." ) except Exception as e: logger.error( f"DownloadCoordinator: Error emitting existing download progress: {e}" ) await anyio.sleep(60) ================================================ FILE: src/exo/download/download_utils.py ================================================ import asyncio import hashlib import os import shutil import ssl import time import traceback from collections.abc import Awaitable from datetime import timedelta from pathlib import Path from typing import Callable, Literal from urllib.parse import urljoin import aiofiles import aiofiles.os as aios import aiohttp import certifi from huggingface_hub import ( snapshot_download, # pyright: ignore[reportUnknownVariableType] ) from loguru import logger from pydantic import ( TypeAdapter, ) from exo.download.huggingface_utils import ( filter_repo_objects, get_allow_patterns, get_auth_headers, get_hf_endpoint, get_hf_token, ) from exo.shared.constants import EXO_MODELS_DIR, EXO_MODELS_PATH from exo.shared.models.model_cards import ModelTask from exo.shared.types.common import ModelId from exo.shared.types.memory import Memory from exo.shared.types.worker.downloads import ( DownloadProgressData, FileListEntry, ModelSafetensorsIndex, RepoDownloadProgress, RepoFileDownloadProgress, ) from exo.shared.types.worker.shards import ShardMetadata class HuggingFaceAuthenticationError(Exception): """Raised when HuggingFace returns 401/403 for a model download.""" class HuggingFaceRateLimitError(Exception): """429 Huggingface code""" async def _build_auth_error_message(status_code: int, model_id: ModelId) -> str: token = await get_hf_token() if status_code == 401 and token is None: return ( f"Model '{model_id}' requires authentication. " f"Set HF_TOKEN in the app's Advanced settings, set the HF_TOKEN environment variable, or run `hf auth login`. " f"Get a token at https://huggingface.co/settings/tokens" ) elif status_code == 403: return ( f"Access denied to '{model_id}'. " f"Please accept the model terms at https://huggingface.co/{model_id}" ) else: return f"Authentication failed for '{model_id}' (HTTP {status_code})" def trim_etag(etag: str) -> str: if (etag[0] == '"' and etag[-1] == '"') or (etag[0] == "'" and etag[-1] == "'"): return etag[1:-1] return etag def map_repo_file_download_progress_to_download_progress_data( repo_file_download_progress: RepoFileDownloadProgress, ) -> DownloadProgressData: return DownloadProgressData( downloaded=repo_file_download_progress.downloaded, downloaded_this_session=repo_file_download_progress.downloaded_this_session, total=repo_file_download_progress.total, completed_files=1 if repo_file_download_progress.status == "complete" else 0, total_files=1, speed=repo_file_download_progress.speed, eta_ms=int(repo_file_download_progress.eta.total_seconds() * 1000), files={}, ) def map_repo_download_progress_to_download_progress_data( repo_download_progress: RepoDownloadProgress, ) -> DownloadProgressData: return DownloadProgressData( total=repo_download_progress.total, downloaded=repo_download_progress.downloaded, downloaded_this_session=repo_download_progress.downloaded_this_session, completed_files=repo_download_progress.completed_files, total_files=repo_download_progress.total_files, speed=repo_download_progress.overall_speed, eta_ms=int(repo_download_progress.overall_eta.total_seconds() * 1000), files={ file_path: map_repo_file_download_progress_to_download_progress_data( file_progress ) for file_path, file_progress in repo_download_progress.file_progress.items() }, ) def resolve_model_in_path(model_id: ModelId) -> Path | None: """Search EXO_MODELS_PATH directories for a pre-existing model. Checks each directory for the normalized name (org--model). A candidate is only returned if ``is_model_directory_complete`` confirms all weight files are present. """ if EXO_MODELS_PATH is None: return None normalized = model_id.normalize() for search_dir in EXO_MODELS_PATH: candidate = search_dir / normalized if candidate.is_dir() and is_model_directory_complete(candidate): return candidate return None def build_model_path(model_id: ModelId) -> Path: found = resolve_model_in_path(model_id) if found is not None: return found return EXO_MODELS_DIR / model_id.normalize() async def resolve_model_path_for_repo(model_id: ModelId) -> Path: return (await ensure_models_dir()) / model_id.normalize() async def ensure_models_dir() -> Path: await aios.makedirs(EXO_MODELS_DIR, exist_ok=True) return EXO_MODELS_DIR async def delete_model(model_id: ModelId) -> bool: models_dir = await ensure_models_dir() model_dir = models_dir / model_id.normalize() cache_dir = models_dir / "caches" / model_id.normalize() deleted = False if await aios.path.exists(model_dir): await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False) deleted = True # Also clear cache if await aios.path.exists(cache_dir): await asyncio.to_thread(shutil.rmtree, cache_dir, ignore_errors=False) return deleted async def seed_models(seed_dir: str | Path): """Move models from resources folder to EXO_MODELS_DIR.""" source_dir = Path(seed_dir) dest_dir = await ensure_models_dir() for path in source_dir.iterdir(): if path.is_dir() and path.name.startswith("models--"): dest_path = dest_dir / path.name if await aios.path.exists(dest_path): logger.info("Skipping moving model to .cache directory") else: try: await aios.rename(str(path), str(dest_path)) except Exception: logger.error(f"Error seeding model {path} to {dest_path}") logger.error(traceback.format_exc()) def _scan_model_directory( model_dir: Path, recursive: bool = False ) -> list[FileListEntry] | None: """Scan a local model directory and build a file list. Requires at least one ``*.safetensors.index.json``. Every weight file referenced by the index that is missing on disk gets ``size=None``. """ index_files = list(model_dir.glob("**/*.safetensors.index.json")) if not index_files: return None entries_by_path: dict[str, FileListEntry] = {} if recursive: for dirpath, _, filenames in os.walk(model_dir): for filename in filenames: if filename.endswith(".partial"): continue full_path = Path(dirpath) / filename rel_path = str(full_path.relative_to(model_dir)) entries_by_path[rel_path] = FileListEntry( type="file", path=rel_path, size=full_path.stat().st_size, ) else: for item in model_dir.iterdir(): if item.is_file() and not item.name.endswith(".partial"): entries_by_path[item.name] = FileListEntry( type="file", path=item.name, size=item.stat().st_size, ) # Add expected weight files from index that haven't been downloaded yet for index_file in index_files: try: index_data = ModelSafetensorsIndex.model_validate_json( index_file.read_text() ) relative_dir = index_file.parent.relative_to(model_dir) for filename in set(index_data.weight_map.values()): rel_path = ( str(relative_dir / filename) if relative_dir != Path(".") else filename ) if rel_path not in entries_by_path: entries_by_path[rel_path] = FileListEntry( type="file", path=rel_path, size=None, ) except Exception: continue return list(entries_by_path.values()) def is_model_directory_complete(model_dir: Path) -> bool: """Check if a model directory contains all required weight files.""" file_list = _scan_model_directory(model_dir, recursive=True) return file_list is not None and all(f.size is not None for f in file_list) async def _build_file_list_from_local_directory( model_id: ModelId, recursive: bool = False, ) -> list[FileListEntry] | None: """Build a file list from locally existing model files. We can only figure out the files we need from safetensors index, so a local directory must contain a *.safetensors.index.json and safetensors listed there. """ model_dir = (await ensure_models_dir()) / model_id.normalize() if not await aios.path.exists(model_dir): return None file_list = await asyncio.to_thread(_scan_model_directory, model_dir, recursive) if not file_list: return None return file_list _fetched_file_lists_this_session: set[str] = set() async def fetch_file_list_with_cache( model_id: ModelId, revision: str = "main", recursive: bool = False, skip_internet: bool = False, on_connection_lost: Callable[[], None] = lambda: None, ) -> list[FileListEntry]: target_dir = (await ensure_models_dir()) / "caches" / model_id.normalize() await aios.makedirs(target_dir, exist_ok=True) cache_file = target_dir / f"{model_id.normalize()}--{revision}--file_list.json" cache_key = f"{model_id.normalize()}--{revision}" if cache_key in _fetched_file_lists_this_session and await aios.path.exists( cache_file ): async with aiofiles.open(cache_file, "r") as f: return TypeAdapter(list[FileListEntry]).validate_json(await f.read()) if skip_internet: if await aios.path.exists(cache_file): async with aiofiles.open(cache_file, "r") as f: return TypeAdapter(list[FileListEntry]).validate_json(await f.read()) local_file_list = await _build_file_list_from_local_directory( model_id, recursive ) if local_file_list is not None: logger.warning( f"No internet and no cached file list for {model_id} - using local file list" ) return local_file_list raise FileNotFoundError( f"No internet connection and no cached file list for {model_id}" ) try: file_list = await fetch_file_list_with_retry( model_id, revision, recursive=recursive, on_connection_lost=on_connection_lost, ) async with aiofiles.open(cache_file, "w") as f: await f.write( TypeAdapter(list[FileListEntry]).dump_json(file_list).decode() ) _fetched_file_lists_this_session.add(cache_key) return file_list except Exception as e: logger.opt(exception=e).warning( "Ran into exception when fetching file list from HF." ) if await aios.path.exists(cache_file): logger.warning( f"No cached file list for {model_id} - using local file list" ) async with aiofiles.open(cache_file, "r") as f: return TypeAdapter(list[FileListEntry]).validate_json(await f.read()) local_file_list = await _build_file_list_from_local_directory( model_id, recursive ) if local_file_list is not None: logger.warning( f"Failed to fetch file list for {model_id} and no cache exists, " ) return local_file_list raise FileNotFoundError(f"Failed to fetch file list for {model_id}: {e}") from e async def fetch_file_list_with_retry( model_id: ModelId, revision: str = "main", path: str = "", recursive: bool = False, on_connection_lost: Callable[[], None] = lambda: None, ) -> list[FileListEntry]: n_attempts = 3 for attempt in range(n_attempts): try: return await _fetch_file_list(model_id, revision, path, recursive) except HuggingFaceAuthenticationError: raise except Exception as e: on_connection_lost() if attempt == n_attempts - 1: raise e await asyncio.sleep(2.0**attempt) raise Exception( f"Failed to fetch file list for {model_id=} {revision=} {path=} {recursive=}" ) async def _fetch_file_list( model_id: ModelId, revision: str = "main", path: str = "", recursive: bool = False ) -> list[FileListEntry]: api_url = f"{get_hf_endpoint()}/api/models/{model_id}/tree/{revision}" url = f"{api_url}/{path}" if path else api_url headers = await get_download_headers() async with ( create_http_session(timeout_profile="short") as session, session.get(url, headers=headers) as response, ): if response.status in [401, 403]: msg = await _build_auth_error_message(response.status, model_id) raise HuggingFaceAuthenticationError(msg) elif response.status == 429: raise HuggingFaceRateLimitError( f"Couldn't download {model_id} because of HuggingFace rate limit." ) elif response.status == 200: data_json = await response.text() data = TypeAdapter(list[FileListEntry]).validate_json(data_json) files: list[FileListEntry] = [] for item in data: if item.type == "file": files.append(FileListEntry.model_validate(item)) elif item.type == "directory" and recursive: subfiles = await _fetch_file_list( model_id, revision, item.path, recursive ) files.extend(subfiles) return files else: raise Exception(f"Failed to fetch file list: {response.status}") async def get_download_headers() -> dict[str, str]: return {**(await get_auth_headers()), "Accept-Encoding": "identity"} def create_http_session( auto_decompress: bool = False, timeout_profile: Literal["short", "long"] = "long", ) -> aiohttp.ClientSession: if timeout_profile == "short": total_timeout = 30 connect_timeout = 10 sock_read_timeout = 30 sock_connect_timeout = 10 else: total_timeout = 1800 connect_timeout = 60 sock_read_timeout = 60 sock_connect_timeout = 60 ssl_context = ssl.create_default_context( cafile=os.getenv("SSL_CERT_FILE") or certifi.where() ) connector = aiohttp.TCPConnector(ssl=ssl_context) return aiohttp.ClientSession( auto_decompress=auto_decompress, connector=connector, proxy=os.getenv("HTTPS_PROXY") or os.getenv("HTTP_PROXY") or None, timeout=aiohttp.ClientTimeout( total=total_timeout, connect=connect_timeout, sock_read=sock_read_timeout, sock_connect=sock_connect_timeout, ), ) async def calc_hash(path: Path, hash_type: Literal["sha1", "sha256"] = "sha1") -> str: hasher = hashlib.sha1() if hash_type == "sha1" else hashlib.sha256() if hash_type == "sha1": header = f"blob {(await aios.stat(path)).st_size}\0".encode() hasher.update(header) async with aiofiles.open(path, "rb") as f: while chunk := await f.read(8 * 1024 * 1024): hasher.update(chunk) return hasher.hexdigest() async def file_meta( model_id: ModelId, revision: str, path: str, redirected_location: str | None = None ) -> tuple[int, str]: url = ( urljoin(f"{get_hf_endpoint()}/{model_id}/resolve/{revision}/", path) if redirected_location is None else f"{get_hf_endpoint()}{redirected_location}" ) headers = await get_download_headers() async with ( create_http_session(timeout_profile="short") as session, session.head(url, headers=headers) as r, ): if r.status == 307: # On redirect, only trust Hugging Face's x-linked-* headers. x_linked_size = r.headers.get("x-linked-size") x_linked_etag = r.headers.get("x-linked-etag") if x_linked_size and x_linked_etag: content_length = int(x_linked_size) etag = trim_etag(x_linked_etag) return content_length, etag # Otherwise, follow the redirect to get authoritative size/hash redirected_location = r.headers.get("location") return await file_meta(model_id, revision, path, redirected_location) if r.status in [401, 403]: msg = await _build_auth_error_message(r.status, model_id) raise HuggingFaceAuthenticationError(msg) content_length = int( r.headers.get("x-linked-size") or r.headers.get("content-length") or 0 ) etag = r.headers.get("x-linked-etag") or r.headers.get("etag") assert content_length > 0, f"No content length for {url}" assert etag is not None, f"No remote hash for {url}" etag = trim_etag(etag) return content_length, etag async def download_file_with_retry( model_id: ModelId, revision: str, path: str, target_dir: Path, on_progress: Callable[[int, int, bool], None] = lambda _, __, ___: None, on_connection_lost: Callable[[], None] = lambda: None, skip_internet: bool = False, ) -> Path: n_attempts = 3 for attempt in range(n_attempts): try: return await _download_file( model_id, revision, path, target_dir, on_progress, skip_internet ) except HuggingFaceAuthenticationError: raise except FileNotFoundError: raise except HuggingFaceRateLimitError as e: if attempt == n_attempts - 1: raise e logger.error( f"Download error on attempt {attempt}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}" ) logger.error(traceback.format_exc()) await asyncio.sleep(2.0**attempt) except Exception as e: if attempt == n_attempts - 1: on_connection_lost() raise e logger.error( f"Download error on attempt {attempt + 1}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}" ) logger.error(traceback.format_exc()) await asyncio.sleep(2.0**attempt) raise Exception( f"Failed to download file {model_id=} {revision=} {path=} {target_dir=}" ) async def _download_file( model_id: ModelId, revision: str, path: str, target_dir: Path, on_progress: Callable[[int, int, bool], None] = lambda _, __, ___: None, skip_internet: bool = False, ) -> Path: target_path = target_dir / path if await aios.path.exists(target_path): if skip_internet: return target_path local_size = (await aios.stat(target_path)).st_size # Try to verify against remote, but allow offline operation try: remote_size, _ = await file_meta(model_id, revision, path) if local_size != remote_size: logger.info( f"File {path} size mismatch (local={local_size}, remote={remote_size}), re-downloading" ) await aios.remove(target_path) else: return target_path except Exception as e: # Offline or network error - trust local file logger.debug( f"Could not verify {path} against remote (offline?): {e}, using local file" ) return target_path if skip_internet: raise FileNotFoundError( f"File {path} not found locally and cannot download in offline mode" ) await aios.makedirs((target_dir / path).parent, exist_ok=True) length, etag = await file_meta(model_id, revision, path) remote_hash = etag[:-5] if etag.endswith("-gzip") else etag partial_path = target_dir / f"{path}.partial" resume_byte_pos = ( (await aios.stat(partial_path)).st_size if (await aios.path.exists(partial_path)) else None ) if resume_byte_pos != length: url = urljoin(f"{get_hf_endpoint()}/{model_id}/resolve/{revision}/", path) headers = await get_download_headers() if resume_byte_pos: headers["Range"] = f"bytes={resume_byte_pos}-" n_read = resume_byte_pos or 0 async with ( create_http_session(timeout_profile="long") as session, session.get(url, headers=headers) as r, ): if r.status == 404: raise FileNotFoundError(f"File not found: {url}") if r.status in [401, 403]: msg = await _build_auth_error_message(r.status, model_id) raise HuggingFaceAuthenticationError(msg) assert r.status in [200, 206], ( f"Failed to download {path} from {url}: {r.status}" ) async with aiofiles.open( partial_path, "ab" if resume_byte_pos else "wb" ) as f: while chunk := await r.content.read(8 * 1024 * 1024): n_read = n_read + (await f.write(chunk)) on_progress(n_read, length, False) final_hash = await calc_hash( partial_path, hash_type="sha256" if len(remote_hash) == 64 else "sha1" ) integrity = final_hash == remote_hash if not integrity: try: await aios.remove(partial_path) except Exception as e: logger.error(f"Error removing partial file {partial_path}: {e}") raise Exception( f"Downloaded file {target_dir / path} has hash {final_hash} but remote hash is {remote_hash}" ) await aios.rename(partial_path, target_dir / path) on_progress(length, length, True) return target_dir / path def calculate_repo_progress( shard: ShardMetadata, model_id: ModelId, revision: str, file_progress: dict[str, RepoFileDownloadProgress], all_start_time: float, ) -> RepoDownloadProgress: all_total = sum((p.total for p in file_progress.values()), Memory.from_bytes(0)) all_downloaded = sum( (p.downloaded for p in file_progress.values()), Memory.from_bytes(0) ) all_downloaded_this_session = sum( (p.downloaded_this_session for p in file_progress.values()), Memory.from_bytes(0), ) elapsed_time = time.time() - all_start_time all_speed = ( all_downloaded_this_session.in_bytes / elapsed_time if elapsed_time > 0 else 0 ) all_eta = ( timedelta(seconds=(all_total - all_downloaded).in_bytes / all_speed) if all_speed > 0 else timedelta(seconds=0) ) status = ( "complete" if all(p.status == "complete" for p in file_progress.values()) else "in_progress" if any(p.status == "in_progress" for p in file_progress.values()) else "not_started" ) return RepoDownloadProgress( repo_id=model_id, repo_revision=revision, shard=shard, completed_files=len( [p for p in file_progress.values() if p.downloaded == p.total] ), total_files=len(file_progress), downloaded=all_downloaded, downloaded_this_session=all_downloaded_this_session, total=all_total, overall_speed=all_speed, overall_eta=all_eta, status=status, file_progress=file_progress, ) async def get_weight_map(model_id: ModelId, revision: str = "main") -> dict[str, str]: target_dir = (await ensure_models_dir()) / model_id.normalize() await aios.makedirs(target_dir, exist_ok=True) index_files_dir = snapshot_download( repo_id=model_id, local_dir=target_dir, allow_patterns="*.safetensors.index.json", ) index_files = list(Path(index_files_dir).glob("**/*.safetensors.index.json")) weight_map: dict[str, str] = {} for index_file in index_files: relative_dir = index_file.parent.relative_to(index_files_dir) async with aiofiles.open(index_file, "r") as f: index_data = ModelSafetensorsIndex.model_validate_json(await f.read()) if relative_dir != Path("."): prefixed_weight_map = { f"{relative_dir}/{key}": str(relative_dir / value) for key, value in index_data.weight_map.items() } weight_map = weight_map | prefixed_weight_map else: weight_map = weight_map | index_data.weight_map return weight_map async def resolve_allow_patterns(shard: ShardMetadata) -> list[str]: # TODO: 'Smart' downloads are disabled because: # (i) We don't handle all kinds of files; # (ii) We don't have sticky sessions. # (iii) Tensor parallel requires all files. return ["*"] try: weight_map = await get_weight_map(str(shard.model_card.model_id)) return get_allow_patterns(weight_map, shard) except Exception: logger.error(f"Error getting weight map for {shard.model_card.model_id=}") logger.error(traceback.format_exc()) return ["*"] def is_image_model(shard: ShardMetadata) -> bool: tasks = shard.model_card.tasks return ModelTask.TextToImage in tasks or ModelTask.ImageToImage in tasks async def get_downloaded_size(path: Path) -> int: partial_path = path.with_suffix(path.suffix + ".partial") if await aios.path.exists(path): return (await aios.stat(path)).st_size if await aios.path.exists(partial_path): return (await aios.stat(partial_path)).st_size return 0 async def download_shard( shard: ShardMetadata, on_progress: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]], max_parallel_downloads: int = 8, skip_download: bool = False, skip_internet: bool = False, allow_patterns: list[str] | None = None, on_connection_lost: Callable[[], None] = lambda: None, ) -> tuple[Path, RepoDownloadProgress]: if not skip_download: logger.debug(f"Downloading {shard.model_card.model_id=}") revision = "main" target_dir = await ensure_models_dir() / str(shard.model_card.model_id).replace( "/", "--" ) if not skip_download: await aios.makedirs(target_dir, exist_ok=True) if not allow_patterns: allow_patterns = await resolve_allow_patterns(shard) if not skip_download: logger.debug(f"Downloading {shard.model_card.model_id=} with {allow_patterns=}") all_start_time = time.time() file_list = await fetch_file_list_with_cache( shard.model_card.model_id, revision, recursive=True, skip_internet=skip_internet, on_connection_lost=on_connection_lost, ) filtered_file_list = list( filter_repo_objects( file_list, allow_patterns=allow_patterns, key=lambda x: x.path ) ) # For image models, skip root-level safetensors files since weights # are stored in component subdirectories (e.g., transformer/, vae/) if is_image_model(shard): filtered_file_list = [ f for f in filtered_file_list if "/" in f.path or not f.path.endswith(".safetensors") ] file_progress: dict[str, RepoFileDownloadProgress] = {} async def on_progress_wrapper( file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool ) -> None: previous_progress = file_progress.get(file.path) # Detect re-download: curr_bytes < previous downloaded means file was deleted and restarted is_redownload = ( previous_progress is not None and curr_bytes < previous_progress.downloaded.in_bytes ) if is_redownload or previous_progress is None: # Fresh download or re-download: reset tracking start_time = time.time() downloaded_this_session = curr_bytes else: # Continuing download: accumulate start_time = previous_progress.start_time downloaded_this_session = ( previous_progress.downloaded_this_session.in_bytes + (curr_bytes - previous_progress.downloaded.in_bytes) ) speed = ( downloaded_this_session / (time.time() - start_time) if time.time() - start_time > 0 else 0 ) eta = ( timedelta(seconds=(total_bytes - curr_bytes) / speed) if speed > 0 else timedelta(seconds=0) ) file_progress[file.path] = RepoFileDownloadProgress( repo_id=shard.model_card.model_id, repo_revision=revision, file_path=file.path, downloaded=Memory.from_bytes(curr_bytes), downloaded_this_session=Memory.from_bytes(downloaded_this_session), total=Memory.from_bytes(total_bytes), speed=speed, eta=eta, status="complete" if curr_bytes == total_bytes and is_renamed else "in_progress", start_time=start_time, ) await on_progress( shard, calculate_repo_progress( shard, shard.model_card.model_id, revision, file_progress, all_start_time, ), ) for file in filtered_file_list: downloaded_bytes = await get_downloaded_size(target_dir / file.path) final_file_exists = await aios.path.exists(target_dir / file.path) file_progress[file.path] = RepoFileDownloadProgress( repo_id=shard.model_card.model_id, repo_revision=revision, file_path=file.path, downloaded=Memory.from_bytes(downloaded_bytes), downloaded_this_session=Memory.from_bytes(0), total=Memory.from_bytes(file.size or 0), speed=0, eta=timedelta(0), status="complete" if final_file_exists and downloaded_bytes == file.size else "not_started", start_time=time.time(), ) semaphore = asyncio.Semaphore(max_parallel_downloads) def schedule_progress( file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool ) -> None: asyncio.create_task( on_progress_wrapper(file, curr_bytes, total_bytes, is_renamed) ) async def download_with_semaphore(file: FileListEntry) -> None: async with semaphore: await download_file_with_retry( shard.model_card.model_id, revision, file.path, target_dir, lambda curr_bytes, total_bytes, is_renamed: schedule_progress( file, curr_bytes, total_bytes, is_renamed ), on_connection_lost=on_connection_lost, skip_internet=skip_internet, ) if not skip_download: await asyncio.gather( *[download_with_semaphore(file) for file in filtered_file_list] ) final_repo_progress = calculate_repo_progress( shard, shard.model_card.model_id, revision, file_progress, all_start_time ) await on_progress(shard, final_repo_progress) if gguf := next((f for f in filtered_file_list if f.path.endswith(".gguf")), None): return target_dir / gguf.path, final_repo_progress else: return target_dir, final_repo_progress ================================================ FILE: src/exo/download/huggingface_utils.py ================================================ import os from fnmatch import fnmatch from pathlib import Path from typing import Callable, Generator, Iterable import aiofiles import aiofiles.os as aios from loguru import logger from exo.shared.types.worker.shards import ShardMetadata def filter_repo_objects[T]( items: Iterable[T], *, allow_patterns: list[str] | str | None = None, ignore_patterns: list[str] | str | None = None, key: Callable[[T], str] | None = None, ) -> Generator[T, None, None]: if isinstance(allow_patterns, str): allow_patterns = [allow_patterns] if isinstance(ignore_patterns, str): ignore_patterns = [ignore_patterns] if allow_patterns is not None: allow_patterns = [_add_wildcard_to_directories(p) for p in allow_patterns] if ignore_patterns is not None: ignore_patterns = [_add_wildcard_to_directories(p) for p in ignore_patterns] if key is None: def _identity(item: T) -> str: if isinstance(item, str): return item if isinstance(item, Path): return str(item) raise ValueError( f"Please provide `key` argument in `filter_repo_objects`: `{item}` is not a string." ) key = _identity for item in items: path = key(item) if allow_patterns is not None and not any( fnmatch(path, r) for r in allow_patterns ): continue if ignore_patterns is not None and any( fnmatch(path, r) for r in ignore_patterns ): continue yield item def _add_wildcard_to_directories(pattern: str) -> str: if pattern[-1] == "/": return pattern + "*" return pattern def get_hf_endpoint() -> str: return os.environ.get("HF_ENDPOINT", "https://huggingface.co") def get_hf_home() -> Path: """Get the Hugging Face home directory.""" return Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface")) async def get_hf_token() -> str | None: """Retrieve the Hugging Face token from HF_TOKEN env var or HF_HOME directory.""" # Check environment variable first if token := os.environ.get("HF_TOKEN"): return token # Fall back to file-based token token_path = get_hf_home() / "token" if await aios.path.exists(token_path): async with aiofiles.open(token_path, "r") as f: return (await f.read()).strip() return None async def get_auth_headers() -> dict[str, str]: """Get authentication headers if a token is available.""" token = await get_hf_token() if token: return {"Authorization": f"Bearer {token}"} return {} def extract_layer_num(tensor_name: str) -> int | None: # This is a simple example and might need to be adjusted based on the actual naming convention parts = tensor_name.split(".") for part in parts: if part.isdigit(): return int(part) return None def get_allow_patterns(weight_map: dict[str, str], shard: ShardMetadata) -> list[str]: default_patterns = set( [ "*.json", "*.py", "tokenizer.model", "tiktoken.model", "*/spiece.model", "*.tiktoken", "*.txt", "*.jinja", ] ) shard_specific_patterns: set[str] = set() if shard.model_card.components is not None: shardable_component = next( (c for c in shard.model_card.components if c.can_shard), None ) if weight_map and shardable_component: for tensor_name, filename in weight_map.items(): # Strip component prefix from tensor name (added by weight map namespacing) # E.g., "transformer/blocks.0.weight" -> "blocks.0.weight" if "/" in tensor_name: _, tensor_name_no_prefix = tensor_name.split("/", 1) else: tensor_name_no_prefix = tensor_name # Determine which component this file belongs to from filename component_path = Path(filename).parts[0] if "/" in filename else None if component_path == shardable_component.component_path.rstrip("/"): layer_num = extract_layer_num(tensor_name_no_prefix) if ( layer_num is not None and shard.start_layer <= layer_num < shard.end_layer ): shard_specific_patterns.add(filename) if shard.is_first_layer or shard.is_last_layer: shard_specific_patterns.add(filename) else: shard_specific_patterns.add(filename) else: shard_specific_patterns = set(["*.safetensors"]) # TODO(ciaran): temporary - Include all files from non-shardable components that have no index file for component in shard.model_card.components: if not component.can_shard and component.safetensors_index_filename is None: component_pattern = f"{component.component_path.rstrip('/')}/*" shard_specific_patterns.add(component_pattern) else: if weight_map: for tensor_name, filename in weight_map.items(): layer_num = extract_layer_num(tensor_name) if ( layer_num is not None and shard.start_layer <= layer_num < shard.end_layer ): shard_specific_patterns.add(filename) layer_independent_files = set( [v for k, v in weight_map.items() if extract_layer_num(k) is None] ) shard_specific_patterns.update(layer_independent_files) logger.debug(f"get_allow_patterns {shard=} {layer_independent_files=}") else: shard_specific_patterns = set(["*.safetensors"]) logger.info(f"get_allow_patterns {shard=} {shard_specific_patterns=}") return list(default_patterns | shard_specific_patterns) ================================================ FILE: src/exo/download/impl_shard_downloader.py ================================================ import asyncio from asyncio import create_task from collections.abc import Awaitable from pathlib import Path from typing import AsyncIterator, Callable from loguru import logger from exo.download.download_utils import RepoDownloadProgress, download_shard from exo.download.shard_downloader import ShardDownloader from exo.shared.models.model_cards import ModelCard, ModelId, get_model_cards from exo.shared.types.worker.shards import ( PipelineShardMetadata, ShardMetadata, ) def exo_shard_downloader( max_parallel_downloads: int = 8, offline: bool = False ) -> ShardDownloader: return SingletonShardDownloader( ResumableShardDownloader(max_parallel_downloads, offline=offline) ) async def build_base_shard(model_id: ModelId) -> ShardMetadata: model_card = await ModelCard.load(model_id) return PipelineShardMetadata( model_card=model_card, device_rank=0, world_size=1, start_layer=0, end_layer=model_card.n_layers, n_layers=model_card.n_layers, ) async def build_full_shard(model_id: ModelId) -> PipelineShardMetadata: base_shard = await build_base_shard(model_id) return PipelineShardMetadata( model_card=base_shard.model_card, device_rank=base_shard.device_rank, world_size=base_shard.world_size, start_layer=base_shard.start_layer, end_layer=base_shard.n_layers, n_layers=base_shard.n_layers, ) class SingletonShardDownloader(ShardDownloader): def __init__(self, shard_downloader: ShardDownloader): self.shard_downloader = shard_downloader self.active_downloads: dict[ShardMetadata, asyncio.Task[Path]] = {} def on_progress( self, callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]], ) -> None: self.shard_downloader.on_progress(callback) async def ensure_shard( self, shard: ShardMetadata, config_only: bool = False ) -> Path: if shard not in self.active_downloads: self.active_downloads[shard] = asyncio.create_task( self.shard_downloader.ensure_shard(shard, config_only) ) try: return await self.active_downloads[shard] finally: if shard in self.active_downloads and self.active_downloads[shard].done(): del self.active_downloads[shard] async def get_shard_download_status( self, ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]: async for path, status in self.shard_downloader.get_shard_download_status(): yield path, status async def get_shard_download_status_for_shard( self, shard: ShardMetadata ) -> RepoDownloadProgress: return await self.shard_downloader.get_shard_download_status_for_shard(shard) class ResumableShardDownloader(ShardDownloader): def __init__(self, max_parallel_downloads: int = 8, offline: bool = False): self.max_parallel_downloads = max_parallel_downloads self.offline = offline self.on_progress_callbacks: list[ Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]] ] = [] async def on_progress_wrapper( self, shard: ShardMetadata, progress: RepoDownloadProgress ) -> None: for callback in self.on_progress_callbacks: await callback(shard, progress) def on_progress( self, callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]], ) -> None: self.on_progress_callbacks.append(callback) async def ensure_shard( self, shard: ShardMetadata, config_only: bool = False ) -> Path: allow_patterns = ["config.json"] if config_only else None target_dir, _ = await download_shard( shard, self.on_progress_wrapper, max_parallel_downloads=self.max_parallel_downloads, allow_patterns=allow_patterns, skip_internet=self.offline, ) return target_dir async def get_shard_download_status( self, ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]: async def _status_for_model( model_id: ModelId, ) -> tuple[Path, RepoDownloadProgress]: """Helper coroutine that builds the shard for a model and gets its download status.""" shard = await build_full_shard(model_id) return await download_shard( shard, self.on_progress_wrapper, skip_download=True, skip_internet=self.offline, ) semaphore = asyncio.Semaphore(self.max_parallel_downloads) async def download_with_semaphore( model_card: ModelCard, ) -> tuple[Path, RepoDownloadProgress]: async with semaphore: return await _status_for_model(model_card.model_id) tasks = [ create_task(download_with_semaphore(model_card)) for model_card in await get_model_cards() ] for task in asyncio.as_completed(tasks): try: yield await task except Exception as e: logger.warning(f"Error downloading shard: {type(e).__name__}") async def get_shard_download_status_for_shard( self, shard: ShardMetadata ) -> RepoDownloadProgress: _, progress = await download_shard( shard, self.on_progress_wrapper, skip_download=True, skip_internet=self.offline, ) return progress ================================================ FILE: src/exo/download/shard_downloader.py ================================================ from abc import ABC, abstractmethod from collections.abc import Awaitable from copy import copy from datetime import timedelta from pathlib import Path from typing import AsyncIterator, Callable from exo.download.download_utils import RepoDownloadProgress from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask from exo.shared.types.memory import Memory from exo.shared.types.worker.shards import ( PipelineShardMetadata, ShardMetadata, ) # TODO: the PipelineShardMetadata getting reinstantiated is a bit messy. Should this be a classmethod? class ShardDownloader(ABC): @abstractmethod async def ensure_shard( self, shard: ShardMetadata, config_only: bool = False ) -> Path: """ Ensures that the shard is downloaded. Does not allow multiple overlapping downloads at once. If you try to download a Shard which overlaps a Shard that is already being downloaded, the download will be cancelled and a new download will start. Args: shard (Shard): The shard to download. """ @abstractmethod def on_progress( self, callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]], ) -> None: pass @abstractmethod async def get_shard_download_status( self, ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]: """Get the download status of shards. Yields: tuple[Path, RepoDownloadProgress]: The path and progress of a shard download. """ yield (Path("/tmp/noop_shard"), NOOP_DOWNLOAD_PROGRESS) @abstractmethod async def get_shard_download_status_for_shard( self, shard: ShardMetadata ) -> RepoDownloadProgress: ... class NoopShardDownloader(ShardDownloader): async def ensure_shard( self, shard: ShardMetadata, config_only: bool = False ) -> Path: return Path("/tmp/noop_shard") def on_progress( self, callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]], ) -> None: pass async def get_shard_download_status( self, ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]: yield ( Path("/tmp/noop_shard"), NOOP_DOWNLOAD_PROGRESS, ) async def get_shard_download_status_for_shard( self, shard: ShardMetadata ) -> RepoDownloadProgress: dp = copy(NOOP_DOWNLOAD_PROGRESS) dp.shard = shard return dp NOOP_DOWNLOAD_PROGRESS = RepoDownloadProgress( repo_id="noop", repo_revision="noop", shard=PipelineShardMetadata( model_card=ModelCard( model_id=ModelId("noop"), storage_size=Memory.from_bytes(0), n_layers=1, hidden_size=1, supports_tensor=False, tasks=[ModelTask.TextGeneration], ), device_rank=0, world_size=1, start_layer=0, end_layer=1, n_layers=1, ), completed_files=0, total_files=0, downloaded=Memory.from_bytes(0), downloaded_this_session=Memory.from_bytes(0), total=Memory.from_bytes(0), overall_speed=0, overall_eta=timedelta(seconds=0), status="complete", ) ================================================ FILE: src/exo/download/tests/__init__.py ================================================ ================================================ FILE: src/exo/download/tests/test_download_verification.py ================================================ """Tests for download verification and cache behavior.""" import time from collections.abc import AsyncIterator from datetime import timedelta from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch import aiofiles import aiofiles.os as aios import pytest from pydantic import TypeAdapter from exo.download.download_utils import ( delete_model, fetch_file_list_with_cache, ) from exo.shared.types.common import ModelId from exo.shared.types.memory import Memory from exo.shared.types.worker.downloads import FileListEntry, RepoFileDownloadProgress @pytest.fixture def model_id() -> ModelId: return ModelId("test-org/test-model") @pytest.fixture async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]: """Set up a temporary models directory for testing.""" models_dir = tmp_path / "models" await aios.makedirs(models_dir, exist_ok=True) with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir): yield models_dir class TestFileVerification: """Tests for file size verification in _download_file.""" async def test_redownload_when_file_size_changes_upstream( self, model_id: ModelId, tmp_path: Path ) -> None: """Test that files with mismatched sizes are re-downloaded.""" # Import inside test to allow patching from exo.download.download_utils import ( _download_file, # pyright: ignore[reportPrivateUsage] ) target_dir = tmp_path / "downloads" await aios.makedirs(target_dir, exist_ok=True) # Create a local file with wrong size local_file = target_dir / "test.safetensors" async with aiofiles.open(local_file, "wb") as f: await f.write(b"local content") # 13 bytes remote_size = 1000 # Different from local remote_hash = "abc123" with ( patch( "exo.download.download_utils.file_meta", new_callable=AsyncMock, return_value=(remote_size, remote_hash), ) as mock_file_meta, patch( "exo.download.download_utils.create_http_session" ) as mock_session_factory, ): # Set up mock HTTP response for re-download mock_response = MagicMock() mock_response.status = 200 mock_response.content.read = AsyncMock( # pyright: ignore[reportAny] side_effect=[b"x" * remote_size, b""] ) mock_session = MagicMock() mock_session.get.return_value.__aenter__ = AsyncMock( # pyright: ignore[reportAny] return_value=mock_response ) mock_session.get.return_value.__aexit__ = AsyncMock( # pyright: ignore[reportAny] return_value=None ) mock_session_factory.return_value.__aenter__ = AsyncMock( # pyright: ignore[reportAny] return_value=mock_session ) mock_session_factory.return_value.__aexit__ = AsyncMock( # pyright: ignore[reportAny] return_value=None ) # Mock calc_hash to return the expected hash with patch( "exo.download.download_utils.calc_hash", new_callable=AsyncMock, return_value=remote_hash, ): await _download_file(model_id, "main", "test.safetensors", target_dir) # file_meta should be called twice: once for verification, once for download assert mock_file_meta.call_count == 2 async def test_skip_download_when_file_size_matches( self, model_id: ModelId, tmp_path: Path ) -> None: """Test that files with matching sizes are not re-downloaded.""" from exo.download.download_utils import ( _download_file, # pyright: ignore[reportPrivateUsage] ) target_dir = tmp_path / "downloads" await aios.makedirs(target_dir, exist_ok=True) # Create a local file local_file = target_dir / "test.safetensors" local_content = b"local content" async with aiofiles.open(local_file, "wb") as f: await f.write(local_content) remote_size = len(local_content) # Same as local remote_hash = "abc123" with ( patch( "exo.download.download_utils.file_meta", new_callable=AsyncMock, return_value=(remote_size, remote_hash), ) as mock_file_meta, patch( "exo.download.download_utils.create_http_session" ) as mock_session_factory, ): result = await _download_file( model_id, "main", "test.safetensors", target_dir ) # Should return immediately without downloading assert result == local_file mock_file_meta.assert_called_once() mock_session_factory.assert_not_called() async def test_offline_fallback_uses_local_file( self, model_id: ModelId, tmp_path: Path ) -> None: """Test that local files are used when network is unavailable.""" from exo.download.download_utils import ( _download_file, # pyright: ignore[reportPrivateUsage] ) target_dir = tmp_path / "downloads" await aios.makedirs(target_dir, exist_ok=True) # Create a local file local_file = target_dir / "test.safetensors" async with aiofiles.open(local_file, "wb") as f: await f.write(b"local content") with ( patch( "exo.download.download_utils.file_meta", new_callable=AsyncMock, side_effect=Exception("Network error"), ), patch( "exo.download.download_utils.create_http_session" ) as mock_session_factory, ): result = await _download_file( model_id, "main", "test.safetensors", target_dir ) # Should return local file without attempting download assert result == local_file mock_session_factory.assert_not_called() class TestFileListCache: """Tests for file list caching behavior.""" async def test_fetch_fresh_and_update_cache( self, model_id: ModelId, tmp_path: Path ) -> None: """Test that fresh data is fetched and cache is updated.""" models_dir = tmp_path / "models" file_list = [ FileListEntry(type="file", path="model.safetensors", size=1000), FileListEntry(type="file", path="config.json", size=100), ] with ( patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir), patch( "exo.download.download_utils.fetch_file_list_with_retry", new_callable=AsyncMock, return_value=file_list, ) as mock_fetch, ): result = await fetch_file_list_with_cache(model_id, "main") assert result == file_list mock_fetch.assert_called_once() # Verify cache was written cache_file = ( models_dir / "caches" / model_id.normalize() / f"{model_id.normalize()}--main--file_list.json" ) assert await aios.path.exists(cache_file) async with aiofiles.open(cache_file, "r") as f: cached_data = TypeAdapter(list[FileListEntry]).validate_json( await f.read() ) assert cached_data == file_list async def test_fallback_to_cache_when_fetch_fails( self, model_id: ModelId, tmp_path: Path ) -> None: """Test that cached data is used when fetch fails.""" models_dir = tmp_path / "models" cache_dir = models_dir / "caches" / model_id.normalize() await aios.makedirs(cache_dir, exist_ok=True) # Create cache file cached_file_list = [ FileListEntry(type="file", path="model.safetensors", size=1000), ] cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json" async with aiofiles.open(cache_file, "w") as f: await f.write( TypeAdapter(list[FileListEntry]).dump_json(cached_file_list).decode() ) with ( patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir), patch( "exo.download.download_utils.fetch_file_list_with_retry", new_callable=AsyncMock, side_effect=Exception("Network error"), ), ): result = await fetch_file_list_with_cache(model_id, "main") assert result == cached_file_list async def test_error_propagates_when_no_cache( self, model_id: ModelId, tmp_path: Path ) -> None: """Test that errors propagate when fetch fails and no cache exists.""" models_dir = tmp_path / "models" with ( patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir), patch( "exo.download.download_utils.fetch_file_list_with_retry", new_callable=AsyncMock, side_effect=Exception("Network error"), ), pytest.raises(Exception, match="Network error"), ): await fetch_file_list_with_cache(model_id, "main") class TestModelDeletion: """Tests for model deletion including cache cleanup.""" async def test_delete_model_clears_cache( self, model_id: ModelId, tmp_path: Path ) -> None: """Test that deleting a model also deletes its cache.""" models_dir = tmp_path / "models" model_dir = models_dir / model_id.normalize() cache_dir = models_dir / "caches" / model_id.normalize() # Create model and cache directories await aios.makedirs(model_dir, exist_ok=True) await aios.makedirs(cache_dir, exist_ok=True) # Add some files async with aiofiles.open(model_dir / "model.safetensors", "w") as f: await f.write("model data") async with aiofiles.open(cache_dir / "file_list.json", "w") as f: await f.write("[]") with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir): result = await delete_model(model_id) assert result is True assert not await aios.path.exists(model_dir) assert not await aios.path.exists(cache_dir) async def test_delete_model_only_cache_exists( self, model_id: ModelId, tmp_path: Path ) -> None: """Test deleting when only cache exists (model already deleted).""" models_dir = tmp_path / "models" cache_dir = models_dir / "caches" / model_id.normalize() # Only create cache directory await aios.makedirs(cache_dir, exist_ok=True) async with aiofiles.open(cache_dir / "file_list.json", "w") as f: await f.write("[]") with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir): result = await delete_model(model_id) # Returns False because model dir didn't exist assert result is False # But cache should still be cleaned up assert not await aios.path.exists(cache_dir) async def test_delete_nonexistent_model( self, model_id: ModelId, tmp_path: Path ) -> None: """Test deleting a model that doesn't exist.""" models_dir = tmp_path / "models" await aios.makedirs(models_dir, exist_ok=True) with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir): result = await delete_model(model_id) assert result is False class TestProgressResetOnRedownload: """Tests for progress tracking when files are re-downloaded.""" async def test_progress_resets_correctly_on_redownload( self, model_id: ModelId ) -> None: """Test that progress tracking resets when a file is re-downloaded. When a file is deleted and re-downloaded (due to size mismatch), the progress tracking should reset rather than calculating negative downloaded_this_session values. """ # Simulate file_progress dict as it exists in download_shard file_progress: dict[str, RepoFileDownloadProgress] = {} # Initialize with old file progress (simulating existing large file) old_file_size = 1_500_000_000 # 1.5 GB file_progress["model.safetensors"] = RepoFileDownloadProgress( repo_id=model_id, repo_revision="main", file_path="model.safetensors", downloaded=Memory.from_bytes(old_file_size), downloaded_this_session=Memory.from_bytes(0), total=Memory.from_bytes(old_file_size), speed=0, eta=timedelta(0), status="not_started", start_time=time.time() - 10, # Started 10 seconds ago ) # Simulate the logic from on_progress_wrapper after re-download starts # This is the exact logic from the fixed on_progress_wrapper curr_bytes = 100_000 # 100 KB - new download just started previous_progress = file_progress.get("model.safetensors") # Detect re-download: curr_bytes < previous downloaded is_redownload = ( previous_progress is not None and curr_bytes < previous_progress.downloaded.in_bytes ) if is_redownload or previous_progress is None: # Fresh download or re-download: reset tracking start_time = time.time() downloaded_this_session = curr_bytes else: # Continuing download: accumulate start_time = previous_progress.start_time downloaded_this_session = ( previous_progress.downloaded_this_session.in_bytes + (curr_bytes - previous_progress.downloaded.in_bytes) ) # Key assertions assert is_redownload is True, "Should detect re-download scenario" assert downloaded_this_session == curr_bytes, ( "downloaded_this_session should equal curr_bytes on re-download" ) assert downloaded_this_session > 0, ( "downloaded_this_session should be positive, not negative" ) # Calculate speed (should be positive) elapsed = time.time() - start_time speed = downloaded_this_session / elapsed if elapsed > 0 else 0 assert speed >= 0, "Speed should be non-negative" async def test_progress_accumulates_on_continuing_download( self, model_id: ModelId ) -> None: """Test that progress accumulates correctly for continuing downloads. When a download continues from where it left off (resume), the progress should accumulate correctly. """ file_progress: dict[str, RepoFileDownloadProgress] = {} # Initialize with partial download progress initial_downloaded = 500_000 # 500 KB already downloaded start_time = time.time() - 5 # Started 5 seconds ago file_progress["model.safetensors"] = RepoFileDownloadProgress( repo_id=model_id, repo_revision="main", file_path="model.safetensors", downloaded=Memory.from_bytes(initial_downloaded), downloaded_this_session=Memory.from_bytes(initial_downloaded), total=Memory.from_bytes(1_000_000), speed=100_000, eta=timedelta(seconds=5), status="in_progress", start_time=start_time, ) # Progress callback with more bytes downloaded curr_bytes = 600_000 # 600 KB - continuing download previous_progress = file_progress.get("model.safetensors") # This is NOT a re-download (curr_bytes > previous downloaded) is_redownload = ( previous_progress is not None and curr_bytes < previous_progress.downloaded.in_bytes ) if is_redownload or previous_progress is None: downloaded_this_session = curr_bytes used_start_time = time.time() else: used_start_time = previous_progress.start_time downloaded_this_session = ( previous_progress.downloaded_this_session.in_bytes + (curr_bytes - previous_progress.downloaded.in_bytes) ) # Key assertions assert is_redownload is False, ( "Should NOT detect re-download for continuing download" ) assert used_start_time == start_time, "Should preserve original start_time" expected_session = initial_downloaded + (curr_bytes - initial_downloaded) assert downloaded_this_session == expected_session, ( f"Should accumulate: {downloaded_this_session} == {expected_session}" ) assert downloaded_this_session == 600_000, ( "downloaded_this_session should equal total downloaded so far" ) ================================================ FILE: src/exo/download/tests/test_offline_mode.py ================================================ """Tests for offline/air-gapped mode.""" from collections.abc import AsyncIterator from pathlib import Path from unittest.mock import AsyncMock, patch import aiofiles import aiofiles.os as aios import pytest from exo.download.download_utils import ( _download_file, # pyright: ignore[reportPrivateUsage] download_file_with_retry, fetch_file_list_with_cache, ) from exo.shared.types.common import ModelId from exo.shared.types.worker.downloads import FileListEntry @pytest.fixture def model_id() -> ModelId: return ModelId("test-org/test-model") @pytest.fixture async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]: models_dir = tmp_path / "models" await aios.makedirs(models_dir, exist_ok=True) with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir): yield models_dir class TestDownloadFileOffline: """Tests for _download_file with skip_internet=True.""" async def test_returns_local_file_without_http_verification( self, model_id: ModelId, tmp_path: Path ) -> None: """When skip_internet=True and file exists locally, return it immediately without making any HTTP calls (no file_meta verification).""" target_dir = tmp_path / "downloads" await aios.makedirs(target_dir, exist_ok=True) local_file = target_dir / "model.safetensors" async with aiofiles.open(local_file, "wb") as f: await f.write(b"model weights data") with patch( "exo.download.download_utils.file_meta", new_callable=AsyncMock, ) as mock_file_meta: result = await _download_file( model_id, "main", "model.safetensors", target_dir, skip_internet=True, ) assert result == local_file mock_file_meta.assert_not_called() async def test_raises_file_not_found_for_missing_file( self, model_id: ModelId, tmp_path: Path ) -> None: """When skip_internet=True and file does NOT exist locally, raise FileNotFoundError instead of attempting download.""" target_dir = tmp_path / "downloads" await aios.makedirs(target_dir, exist_ok=True) with pytest.raises(FileNotFoundError, match="offline mode"): await _download_file( model_id, "main", "missing_model.safetensors", target_dir, skip_internet=True, ) async def test_returns_local_file_in_subdirectory( self, model_id: ModelId, tmp_path: Path ) -> None: """When skip_internet=True and file exists in a subdirectory, return it without HTTP calls.""" target_dir = tmp_path / "downloads" subdir = target_dir / "transformer" await aios.makedirs(subdir, exist_ok=True) local_file = subdir / "diffusion_pytorch_model.safetensors" async with aiofiles.open(local_file, "wb") as f: await f.write(b"weights") with patch( "exo.download.download_utils.file_meta", new_callable=AsyncMock, ) as mock_file_meta: result = await _download_file( model_id, "main", "transformer/diffusion_pytorch_model.safetensors", target_dir, skip_internet=True, ) assert result == local_file mock_file_meta.assert_not_called() class TestDownloadFileWithRetryOffline: """Tests for download_file_with_retry with skip_internet=True.""" async def test_propagates_skip_internet_to_download_file( self, model_id: ModelId, tmp_path: Path ) -> None: """Verify skip_internet is passed through to _download_file.""" target_dir = tmp_path / "downloads" await aios.makedirs(target_dir, exist_ok=True) local_file = target_dir / "config.json" async with aiofiles.open(local_file, "wb") as f: await f.write(b'{"model_type": "qwen2"}') with patch( "exo.download.download_utils.file_meta", new_callable=AsyncMock, ) as mock_file_meta: result = await download_file_with_retry( model_id, "main", "config.json", target_dir, skip_internet=True, ) assert result == local_file mock_file_meta.assert_not_called() async def test_file_not_found_does_not_retry( self, model_id: ModelId, tmp_path: Path ) -> None: """FileNotFoundError from offline mode should not trigger retries.""" target_dir = tmp_path / "downloads" await aios.makedirs(target_dir, exist_ok=True) with pytest.raises(FileNotFoundError): await download_file_with_retry( model_id, "main", "nonexistent.safetensors", target_dir, skip_internet=True, ) class TestFetchFileListOffline: """Tests for fetch_file_list_with_cache with skip_internet=True.""" async def test_uses_cached_file_list( self, model_id: ModelId, temp_models_dir: Path ) -> None: """When skip_internet=True and cache file exists, use it without network.""" from pydantic import TypeAdapter cache_dir = temp_models_dir / "caches" / model_id.normalize() await aios.makedirs(cache_dir, exist_ok=True) cached_list = [ FileListEntry(type="file", path="model.safetensors", size=1000), FileListEntry(type="file", path="config.json", size=200), ] cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json" async with aiofiles.open(cache_file, "w") as f: await f.write( TypeAdapter(list[FileListEntry]).dump_json(cached_list).decode() ) with patch( "exo.download.download_utils.fetch_file_list_with_retry", new_callable=AsyncMock, ) as mock_fetch: result = await fetch_file_list_with_cache( model_id, "main", skip_internet=True ) assert result == cached_list mock_fetch.assert_not_called() async def test_falls_back_to_local_directory_scan( self, model_id: ModelId, temp_models_dir: Path ) -> None: """When skip_internet=True and no cache but local files exist, build file list from local directory.""" import json model_dir = temp_models_dir / model_id.normalize() await aios.makedirs(model_dir, exist_ok=True) async with aiofiles.open(model_dir / "config.json", "w") as f: await f.write('{"model_type": "qwen2"}') index_data = { "metadata": {}, "weight_map": {"model.layers.0.weight": "model.safetensors"}, } async with aiofiles.open(model_dir / "model.safetensors.index.json", "w") as f: await f.write(json.dumps(index_data)) async with aiofiles.open(model_dir / "model.safetensors", "wb") as f: await f.write(b"x" * 500) with patch( "exo.download.download_utils.fetch_file_list_with_retry", new_callable=AsyncMock, ) as mock_fetch: result = await fetch_file_list_with_cache( model_id, "main", skip_internet=True ) mock_fetch.assert_not_called() paths = {entry.path for entry in result} assert "config.json" in paths assert "model.safetensors" in paths async def test_raises_when_no_cache_and_no_local_files( self, model_id: ModelId, temp_models_dir: Path ) -> None: """When skip_internet=True and neither cache nor local files exist, raise FileNotFoundError.""" with pytest.raises(FileNotFoundError, match="No internet"): await fetch_file_list_with_cache(model_id, "main", skip_internet=True) ================================================ FILE: src/exo/download/tests/test_re_download.py ================================================ """Tests that re-downloading a previously deleted model completes successfully.""" import asyncio import contextlib from collections.abc import AsyncIterator, Awaitable from datetime import timedelta from pathlib import Path from typing import Callable from unittest.mock import AsyncMock, patch from exo.download.coordinator import DownloadCoordinator from exo.download.download_utils import RepoDownloadProgress from exo.download.impl_shard_downloader import SingletonShardDownloader from exo.download.shard_downloader import ShardDownloader from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask from exo.shared.types.commands import ( DeleteDownload, ForwarderDownloadCommand, StartDownload, ) from exo.shared.types.common import NodeId, SystemId from exo.shared.types.events import Event, NodeDownloadProgress from exo.shared.types.memory import Memory from exo.shared.types.worker.downloads import DownloadCompleted from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata from exo.utils.channels import Receiver, Sender, channel NODE_ID = NodeId("aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa") MODEL_ID = ModelId("test-org/test-model") def _make_shard(model_id: ModelId = MODEL_ID) -> ShardMetadata: return PipelineShardMetadata( model_card=ModelCard( model_id=model_id, storage_size=Memory.from_mb(100), n_layers=28, hidden_size=1024, supports_tensor=False, tasks=[ModelTask.TextGeneration], ), device_rank=0, world_size=1, start_layer=0, end_layer=28, n_layers=28, ) class FakeShardDownloader(ShardDownloader): """Fake downloader that simulates a successful download by firing the progress callback with status='complete' when ensure_shard is called.""" def __init__(self) -> None: self._progress_callbacks: list[ Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]] ] = [] def on_progress( self, callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]], ) -> None: self._progress_callbacks.append(callback) async def ensure_shard( self, shard: ShardMetadata, config_only: bool = False, # noqa: ARG002 ) -> Path: # Simulate a completed download by firing the progress callback progress = RepoDownloadProgress( repo_id=str(shard.model_card.model_id), repo_revision="main", shard=shard, completed_files=1, total_files=1, downloaded=Memory.from_mb(100), downloaded_this_session=Memory.from_mb(100), total=Memory.from_mb(100), overall_speed=0, overall_eta=timedelta(seconds=0), status="complete", ) for cb in self._progress_callbacks: await cb(shard, progress) return Path("/fake/models") / shard.model_card.model_id.normalize() async def get_shard_download_status( self, ) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]: if False: # noqa: SIM108 # empty async generator yield ( Path(), RepoDownloadProgress( # pyright: ignore[reportUnreachable] repo_id="", repo_revision="", shard=_make_shard(), completed_files=0, total_files=0, downloaded=Memory.from_bytes(0), downloaded_this_session=Memory.from_bytes(0), total=Memory.from_bytes(0), overall_speed=0, overall_eta=timedelta(seconds=0), status="not_started", ), ) async def get_shard_download_status_for_shard( self, shard: ShardMetadata, ) -> RepoDownloadProgress: return RepoDownloadProgress( repo_id=str(shard.model_card.model_id), repo_revision="main", shard=shard, completed_files=0, total_files=1, downloaded=Memory.from_bytes(0), downloaded_this_session=Memory.from_bytes(0), total=Memory.from_mb(100), overall_speed=0, overall_eta=timedelta(seconds=0), status="not_started", ) async def test_re_download_after_delete_completes() -> None: """A model that was downloaded, deleted, and then re-downloaded should reach DownloadCompleted status. This is an end-to-end test through the DownloadCoordinator.""" cmd_send: Sender[ForwarderDownloadCommand] cmd_send, cmd_recv = channel[ForwarderDownloadCommand]() event_send, event_recv = channel[Event]() fake_downloader = FakeShardDownloader() wrapped_downloader = SingletonShardDownloader(fake_downloader) coordinator = DownloadCoordinator( node_id=NODE_ID, shard_downloader=wrapped_downloader, download_command_receiver=cmd_recv, event_sender=event_send, ) shard = _make_shard() origin = SystemId("test") with patch("exo.download.coordinator.delete_model", new_callable=AsyncMock): # Run the coordinator in the background coordinator_task = asyncio.create_task(coordinator.run()) try: # 1. Start first download await cmd_send.send( ForwarderDownloadCommand( origin=origin, command=StartDownload(target_node_id=NODE_ID, shard_metadata=shard), ) ) # Wait for DownloadCompleted first_completed = await _wait_for_download_completed(event_recv, MODEL_ID) assert first_completed is not None, "First download should complete" # 2. Delete the model await cmd_send.send( ForwarderDownloadCommand( origin=origin, command=DeleteDownload(target_node_id=NODE_ID, model_id=MODEL_ID), ) ) # Give the coordinator time to process the delete await asyncio.sleep(0.05) # 3. Re-download the same model await cmd_send.send( ForwarderDownloadCommand( origin=origin, command=StartDownload(target_node_id=NODE_ID, shard_metadata=shard), ) ) # Wait for second DownloadCompleted — this is the bug: it never arrives second_completed = await _wait_for_download_completed(event_recv, MODEL_ID) assert second_completed is not None, ( "Re-download after deletion should complete" ) finally: coordinator.shutdown() coordinator_task.cancel() with contextlib.suppress(asyncio.CancelledError): await coordinator_task async def _wait_for_download_completed( event_recv: Receiver[Event], model_id: ModelId, timeout: float = 2.0 ) -> DownloadCompleted | None: """Drain events until we see a DownloadCompleted for the given model, or timeout.""" try: async with asyncio.timeout(timeout): while True: event = await event_recv.receive() if ( isinstance(event, NodeDownloadProgress) and isinstance(event.download_progress, DownloadCompleted) and event.download_progress.shard_metadata.model_card.model_id == model_id ): return event.download_progress except TimeoutError: return None ================================================ FILE: src/exo/main.py ================================================ import argparse import multiprocessing as mp import os import resource import signal from dataclasses import dataclass, field from typing import Self import anyio from loguru import logger from pydantic import PositiveInt import exo.routing.topics as topics from exo.api.main import API from exo.download.coordinator import DownloadCoordinator from exo.download.impl_shard_downloader import exo_shard_downloader from exo.master.main import Master from exo.routing.event_router import EventRouter from exo.routing.router import Router, get_node_id_keypair from exo.shared.constants import EXO_LOG from exo.shared.election import Election, ElectionResult from exo.shared.logging import logger_cleanup, logger_setup from exo.shared.types.common import NodeId, SessionId from exo.utils.channels import Receiver, channel from exo.utils.pydantic_ext import CamelCaseModel from exo.utils.task_group import TaskGroup from exo.worker.main import Worker @dataclass class Node: router: Router event_router: EventRouter download_coordinator: DownloadCoordinator | None worker: Worker | None election: Election # Every node participates in election, as we do want a node to become master even if it isn't a master candidate if no master candidates are present. election_result_receiver: Receiver[ElectionResult] master: Master | None api: API | None node_id: NodeId offline: bool _tg: TaskGroup = field(init=False, default_factory=TaskGroup) @classmethod async def create(cls, args: "Args") -> Self: keypair = get_node_id_keypair() node_id = NodeId(keypair.to_node_id()) session_id = SessionId(master_node_id=node_id, election_clock=0) router = Router.create(keypair) await router.register_topic(topics.GLOBAL_EVENTS) await router.register_topic(topics.LOCAL_EVENTS) await router.register_topic(topics.COMMANDS) await router.register_topic(topics.ELECTION_MESSAGES) await router.register_topic(topics.CONNECTION_MESSAGES) await router.register_topic(topics.DOWNLOAD_COMMANDS) event_router = EventRouter( session_id, command_sender=router.sender(topics.COMMANDS), external_outbound=router.sender(topics.LOCAL_EVENTS), external_inbound=router.receiver(topics.GLOBAL_EVENTS), ) logger.info(f"Starting node {node_id}") # Create DownloadCoordinator (unless --no-downloads) if not args.no_downloads: download_coordinator = DownloadCoordinator( node_id, exo_shard_downloader(offline=args.offline), event_sender=event_router.sender(), download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS), offline=args.offline, ) else: download_coordinator = None if args.spawn_api: api = API( node_id, port=args.api_port, event_receiver=event_router.receiver(), command_sender=router.sender(topics.COMMANDS), download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS), election_receiver=router.receiver(topics.ELECTION_MESSAGES), ) else: api = None if not args.no_worker: worker = Worker( node_id, event_receiver=event_router.receiver(), event_sender=event_router.sender(), command_sender=router.sender(topics.COMMANDS), download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS), ) else: worker = None # We start every node with a master master = Master( node_id, session_id, event_sender=event_router.sender(), global_event_sender=router.sender(topics.GLOBAL_EVENTS), local_event_receiver=router.receiver(topics.LOCAL_EVENTS), command_receiver=router.receiver(topics.COMMANDS), download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS), ) er_send, er_recv = channel[ElectionResult]() election = Election( node_id, # If someone manages to assemble 1 MILLION devices into an exo cluster then. well done. good job champ. seniority=1_000_000 if args.force_master else 0, # nb: this DOES feedback right now. i have thoughts on how to address this, # but ultimately it seems not worth the complexity election_message_sender=router.sender(topics.ELECTION_MESSAGES), election_message_receiver=router.receiver(topics.ELECTION_MESSAGES), connection_message_receiver=router.receiver(topics.CONNECTION_MESSAGES), command_receiver=router.receiver(topics.COMMANDS), election_result_sender=er_send, ) return cls( router, event_router, download_coordinator, worker, election, er_recv, master, api, node_id, args.offline, ) async def run(self): async with self._tg as tg: signal.signal(signal.SIGINT, lambda _, __: self.shutdown()) signal.signal(signal.SIGTERM, lambda _, __: self.shutdown()) tg.start_soon(self.router.run) tg.start_soon(self.event_router.run) tg.start_soon(self.election.run) if self.download_coordinator: tg.start_soon(self.download_coordinator.run) if self.worker: tg.start_soon(self.worker.run) if self.master: tg.start_soon(self.master.run) if self.api: tg.start_soon(self.api.run) tg.start_soon(self._elect_loop) def shutdown(self): # if this is our second call to shutdown, just sys.exit if self._tg.cancel_called(): import sys sys.exit(1) self._tg.cancel_tasks() async def _elect_loop(self): with self.election_result_receiver as results: async for result in results: # This function continues to have a lot of very specific entangled logic # At least it's somewhat contained # I don't like this duplication, but it's manageable for now. # TODO: This function needs refactoring generally # Ok: # On new master: # - Elect master locally if necessary # - Shutdown and re-create the worker # - Shut down and re-create the API if result.is_new_master: await anyio.sleep(0) self.event_router.shutdown() self.event_router = EventRouter( result.session_id, self.router.sender(topics.COMMANDS), self.router.receiver(topics.GLOBAL_EVENTS), self.router.sender(topics.LOCAL_EVENTS), ) self._tg.start_soon(self.event_router.run) if ( result.session_id.master_node_id == self.node_id and self.master is not None ): logger.info("Node elected Master") elif ( result.session_id.master_node_id == self.node_id and self.master is None ): logger.info("Node elected Master - promoting self") self.master = Master( self.node_id, result.session_id, event_sender=self.event_router.sender(), global_event_sender=self.router.sender(topics.GLOBAL_EVENTS), local_event_receiver=self.router.receiver(topics.LOCAL_EVENTS), command_receiver=self.router.receiver(topics.COMMANDS), download_command_sender=self.router.sender( topics.DOWNLOAD_COMMANDS ), ) self._tg.start_soon(self.master.run) elif ( result.session_id.master_node_id != self.node_id and self.master is not None ): logger.info( f"Node {result.session_id.master_node_id} elected master - demoting self" ) await self.master.shutdown() self.master = None else: logger.info( f"Node {result.session_id.master_node_id} elected master" ) if result.is_new_master: if self.download_coordinator: self.download_coordinator.shutdown() self.download_coordinator = DownloadCoordinator( self.node_id, exo_shard_downloader(offline=self.offline), event_sender=self.event_router.sender(), download_command_receiver=self.router.receiver( topics.DOWNLOAD_COMMANDS ), offline=self.offline, ) self._tg.start_soon(self.download_coordinator.run) if self.worker: self.worker.shutdown() # TODO: add profiling etc to resource monitor self.worker = Worker( self.node_id, event_receiver=self.event_router.receiver(), event_sender=self.event_router.sender(), command_sender=self.router.sender(topics.COMMANDS), download_command_sender=self.router.sender( topics.DOWNLOAD_COMMANDS ), ) self._tg.start_soon(self.worker.run) if self.api: self.api.reset(result.won_clock, self.event_router.receiver()) else: if self.api: self.api.unpause(result.won_clock) def main(): args = Args.parse() soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) target = min(max(soft, 65535), hard) resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard)) mp.set_start_method("spawn", force=True) # TODO: Refactor the current verbosity system logger_setup(EXO_LOG, args.verbosity) logger.info("Starting EXO") logger.info(f"EXO_LIBP2P_NAMESPACE: {os.getenv('EXO_LIBP2P_NAMESPACE')}") if args.offline: logger.info("Running in OFFLINE mode — no internet checks, local models only") if args.no_batch: os.environ["EXO_NO_BATCH"] = "1" logger.info("Continuous batching disabled (--no-batch)") # Set FAST_SYNCH override env var for runner subprocesses if args.fast_synch is True: os.environ["EXO_FAST_SYNCH"] = "on" logger.info("FAST_SYNCH forced ON") elif args.fast_synch is False: os.environ["EXO_FAST_SYNCH"] = "off" logger.info("FAST_SYNCH forced OFF") node = anyio.run(Node.create, args) try: anyio.run(node.run) except BaseException as exception: logger.opt(exception=exception).critical( "EXO terminated due to unhandled exception" ) raise finally: logger.info("EXO Shutdown complete") logger_cleanup() class Args(CamelCaseModel): verbosity: int = 0 force_master: bool = False spawn_api: bool = False api_port: PositiveInt = 52415 tb_only: bool = False no_worker: bool = False no_downloads: bool = False offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true" no_batch: bool = False fast_synch: bool | None = None # None = auto, True = force on, False = force off @classmethod def parse(cls) -> Self: parser = argparse.ArgumentParser(prog="EXO") default_verbosity = 0 parser.add_argument( "-q", "--quiet", action="store_const", const=-1, dest="verbosity", default=default_verbosity, ) parser.add_argument( "-v", "--verbose", action="count", dest="verbosity", default=default_verbosity, ) parser.add_argument( "-m", "--force-master", action="store_true", dest="force_master", ) parser.add_argument( "--no-api", action="store_false", dest="spawn_api", ) parser.add_argument( "--api-port", type=int, dest="api_port", default=52415, ) parser.add_argument( "--no-worker", action="store_true", ) parser.add_argument( "--no-downloads", action="store_true", help="Disable the download coordinator (node won't download models)", ) parser.add_argument( "--offline", action="store_true", default=os.getenv("EXO_OFFLINE", "false").lower() == "true", help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models", ) parser.add_argument( "--no-batch", action="store_true", help="Disable continuous batching, use sequential generation", ) fast_synch_group = parser.add_mutually_exclusive_group() fast_synch_group.add_argument( "--fast-synch", action="store_true", dest="fast_synch", default=None, help="Force MLX FAST_SYNCH on (for JACCL backend)", ) fast_synch_group.add_argument( "--no-fast-synch", action="store_false", dest="fast_synch", help="Force MLX FAST_SYNCH off", ) args = parser.parse_args() return cls(**vars(args)) # pyright: ignore[reportAny] - We are intentionally validating here, we can't do it statically ================================================ FILE: src/exo/master/__init__.py ================================================ ================================================ FILE: src/exo/master/image_store.py ================================================ import time from pathlib import Path from pydantic import BaseModel from exo.shared.types.common import Id class StoredImage(BaseModel, frozen=True): image_id: Id file_path: Path content_type: str expires_at: float class ImageStore: def __init__(self, storage_dir: Path, default_expiry_seconds: int = 3600) -> None: self._storage_dir = storage_dir self._default_expiry_seconds = default_expiry_seconds self._images: dict[Id, StoredImage] = {} self._storage_dir.mkdir(parents=True, exist_ok=True) def store(self, image_bytes: bytes, content_type: str) -> StoredImage: image_id = Id() extension = _content_type_to_extension(content_type) file_path = self._storage_dir / f"{image_id}{extension}" file_path.write_bytes(image_bytes) stored = StoredImage( image_id=image_id, file_path=file_path, content_type=content_type, expires_at=time.time() + self._default_expiry_seconds, ) self._images[image_id] = stored return stored def get(self, image_id: Id) -> StoredImage | None: stored = self._images.get(image_id) if stored is None: return None if time.time() > stored.expires_at: self._remove(image_id) return None return stored def list_images(self) -> list[StoredImage]: now = time.time() return [stored for stored in self._images.values() if now <= stored.expires_at] def cleanup_expired(self) -> int: now = time.time() expired_ids = [ image_id for image_id, stored in self._images.items() if now > stored.expires_at ] for image_id in expired_ids: self._remove(image_id) return len(expired_ids) def _remove(self, image_id: Id) -> None: stored = self._images.pop(image_id, None) if stored is not None and stored.file_path.exists(): stored.file_path.unlink() def _content_type_to_extension( content_type: str, ) -> str: ext = f"{content_type.split('/')[1]}" if ext == "jpeg": ext = "jpg" return f".{ext}" ================================================ FILE: src/exo/master/main.py ================================================ from datetime import datetime, timedelta, timezone import anyio from loguru import logger from exo.master.placement import ( add_instance_to_placements, cancel_unnecessary_downloads, delete_instance, get_transition_events, place_instance, ) from exo.shared.apply import apply from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED from exo.shared.types.commands import ( CreateInstance, DeleteInstance, ForwarderCommand, ForwarderDownloadCommand, ImageEdits, ImageGeneration, PlaceInstance, RequestEventLog, SendInputChunk, TaskCancelled, TaskFinished, TestCommand, TextGeneration, ) from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId from exo.shared.types.events import ( Event, GlobalForwarderEvent, IndexedEvent, InputChunkReceived, InstanceDeleted, LocalForwarderEvent, NodeGatheredInfo, NodeTimedOut, TaskCreated, TaskDeleted, TaskStatusUpdated, TraceEventData, TracesCollected, TracesMerged, ) from exo.shared.types.state import State from exo.shared.types.tasks import ( ImageEdits as ImageEditsTask, ) from exo.shared.types.tasks import ( ImageGeneration as ImageGenerationTask, ) from exo.shared.types.tasks import ( TaskId, TaskStatus, ) from exo.shared.types.tasks import ( TextGeneration as TextGenerationTask, ) from exo.shared.types.worker.instances import InstanceId from exo.utils.channels import Receiver, Sender from exo.utils.disk_event_log import DiskEventLog from exo.utils.event_buffer import MultiSourceBuffer from exo.utils.task_group import TaskGroup class Master: def __init__( self, node_id: NodeId, session_id: SessionId, *, command_receiver: Receiver[ForwarderCommand], event_sender: Sender[Event], local_event_receiver: Receiver[LocalForwarderEvent], global_event_sender: Sender[GlobalForwarderEvent], download_command_sender: Sender[ForwarderDownloadCommand], ): self.node_id = node_id self.session_id = session_id self.state = State() self._tg: TaskGroup = TaskGroup() self.command_task_mapping: dict[CommandId, TaskId] = {} self.command_receiver = command_receiver self.local_event_receiver = local_event_receiver self.global_event_sender = global_event_sender self.download_command_sender = download_command_sender self.event_sender = event_sender self._system_id = SystemId() self._multi_buffer = MultiSourceBuffer[SystemId, Event]() self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master") self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {} self._expected_ranks: dict[TaskId, set[int]] = {} async def run(self): logger.info("Starting Master") try: async with self._tg as tg: tg.start_soon(self._event_processor) tg.start_soon(self._command_processor) tg.start_soon(self._plan) finally: self._event_log.close() self.global_event_sender.close() self.local_event_receiver.close() self.command_receiver.close() async def shutdown(self): logger.info("Stopping Master") self._tg.cancel_tasks() async def _command_processor(self) -> None: with self.command_receiver as commands: async for forwarder_command in commands: try: logger.info(f"Executing command: {forwarder_command.command}") generated_events: list[Event] = [] command = forwarder_command.command instance_task_counts: dict[InstanceId, int] = {} match command: case TestCommand(): pass case TextGeneration(): for instance in self.state.instances.values(): if ( instance.shard_assignments.model_id == command.task_params.model ): task_count = sum( 1 for task in self.state.tasks.values() if task.instance_id == instance.instance_id ) instance_task_counts[instance.instance_id] = ( task_count ) if not instance_task_counts: raise ValueError( f"No instance found for model {command.task_params.model}" ) available_instance_ids = sorted( instance_task_counts.keys(), key=lambda instance_id: instance_task_counts[ instance_id ], ) task_id = TaskId() generated_events.append( TaskCreated( task_id=task_id, task=TextGenerationTask( task_id=task_id, command_id=command.command_id, instance_id=available_instance_ids[0], task_status=TaskStatus.Pending, task_params=command.task_params, ), ) ) self.command_task_mapping[command.command_id] = task_id case ImageGeneration(): for instance in self.state.instances.values(): if ( instance.shard_assignments.model_id == command.task_params.model ): task_count = sum( 1 for task in self.state.tasks.values() if task.instance_id == instance.instance_id ) instance_task_counts[instance.instance_id] = ( task_count ) if not instance_task_counts: raise ValueError( f"No instance found for model {command.task_params.model}" ) available_instance_ids = sorted( instance_task_counts.keys(), key=lambda instance_id: instance_task_counts[ instance_id ], ) task_id = TaskId() selected_instance_id = available_instance_ids[0] generated_events.append( TaskCreated( task_id=task_id, task=ImageGenerationTask( task_id=task_id, command_id=command.command_id, instance_id=selected_instance_id, task_status=TaskStatus.Pending, task_params=command.task_params, ), ) ) self.command_task_mapping[command.command_id] = task_id if EXO_TRACING_ENABLED: selected_instance = self.state.instances.get( selected_instance_id ) if selected_instance: ranks = set( shard.device_rank for shard in selected_instance.shard_assignments.runner_to_shard.values() ) self._expected_ranks[task_id] = ranks case ImageEdits(): for instance in self.state.instances.values(): if ( instance.shard_assignments.model_id == command.task_params.model ): task_count = sum( 1 for task in self.state.tasks.values() if task.instance_id == instance.instance_id ) instance_task_counts[instance.instance_id] = ( task_count ) if not instance_task_counts: raise ValueError( f"No instance found for model {command.task_params.model}" ) available_instance_ids = sorted( instance_task_counts.keys(), key=lambda instance_id: instance_task_counts[ instance_id ], ) task_id = TaskId() selected_instance_id = available_instance_ids[0] generated_events.append( TaskCreated( task_id=task_id, task=ImageEditsTask( task_id=task_id, command_id=command.command_id, instance_id=selected_instance_id, task_status=TaskStatus.Pending, task_params=command.task_params, ), ) ) self.command_task_mapping[command.command_id] = task_id if EXO_TRACING_ENABLED: selected_instance = self.state.instances.get( selected_instance_id ) if selected_instance: ranks = set( shard.device_rank for shard in selected_instance.shard_assignments.runner_to_shard.values() ) self._expected_ranks[task_id] = ranks case DeleteInstance(): placement = delete_instance(command, self.state.instances) transition_events = get_transition_events( self.state.instances, placement, self.state.tasks ) for cmd in cancel_unnecessary_downloads( placement, self.state.downloads ): await self.download_command_sender.send( ForwarderDownloadCommand( origin=self._system_id, command=cmd ) ) generated_events.extend(transition_events) case PlaceInstance(): placement = place_instance( command, self.state.topology, self.state.instances, self.state.node_memory, self.state.node_network, ) transition_events = get_transition_events( self.state.instances, placement, self.state.tasks ) generated_events.extend(transition_events) case CreateInstance(): placement = add_instance_to_placements( command, self.state.topology, self.state.instances, ) transition_events = get_transition_events( self.state.instances, placement, self.state.tasks ) generated_events.extend(transition_events) case SendInputChunk(chunk=chunk): generated_events.append( InputChunkReceived( command_id=chunk.command_id, chunk=chunk, ) ) case TaskCancelled(): if ( task_id := self.command_task_mapping.get( command.cancelled_command_id ) ) is not None: generated_events.append( TaskStatusUpdated( task_status=TaskStatus.Cancelled, task_id=task_id, ) ) else: logger.warning( f"Nonexistent command {command.cancelled_command_id} cancelled" ) case TaskFinished(): if ( task_id := self.command_task_mapping.pop( command.finished_command_id, None ) ) is not None: generated_events.append(TaskDeleted(task_id=task_id)) else: logger.warning( f"Finished command {command.finished_command_id} finished" ) case RequestEventLog(): # We should just be able to send everything, since other buffers will ignore old messages # rate limit to 1000 at a time end = min(command.since_idx + 1000, len(self._event_log)) for i, event in enumerate( self._event_log.read_range(command.since_idx, end), start=command.since_idx, ): await self._send_event(IndexedEvent(idx=i, event=event)) for event in generated_events: await self.event_sender.send(event) except ValueError as e: logger.opt(exception=e).warning("Error in command processor") # These plan loops are the cracks showing in our event sourcing architecture - more things could be commands async def _plan(self) -> None: while True: # kill broken instances connected_node_ids = set(self.state.topology.list_nodes()) for instance_id, instance in self.state.instances.items(): for node_id in instance.shard_assignments.node_to_runner: if node_id not in connected_node_ids: await self.event_sender.send( InstanceDeleted(instance_id=instance_id) ) break # time out dead nodes for node_id, time in self.state.last_seen.items(): now = datetime.now(tz=timezone.utc) if now - time > timedelta(seconds=30): logger.info(f"Manually removing node {node_id} due to inactivity") await self.event_sender.send(NodeTimedOut(node_id=node_id)) await anyio.sleep(10) async def _event_processor(self) -> None: with self.local_event_receiver as local_events: async for local_event in local_events: # Discard all events not from our session if local_event.session != self.session_id: continue self._multi_buffer.ingest( local_event.origin_idx, local_event.event, local_event.origin, ) for event in self._multi_buffer.drain(): if isinstance(event, TracesCollected): await self._handle_traces_collected(event) continue logger.debug(f"Master indexing event: {str(event)[:100]}") indexed = IndexedEvent(event=event, idx=len(self._event_log)) self.state = apply(self.state, indexed) event._master_time_stamp = datetime.now(tz=timezone.utc) # pyright: ignore[reportPrivateUsage] if isinstance(event, NodeGatheredInfo): event.when = str(datetime.now(tz=timezone.utc)) self._event_log.append(event) await self._send_event(indexed) # This function is re-entrant, take care! async def _send_event(self, event: IndexedEvent): # Convenience method since this line is ugly await self.global_event_sender.send( GlobalForwarderEvent( origin=self.node_id, origin_idx=event.idx, session=self.session_id, event=event.event, ) ) async def _handle_traces_collected(self, event: TracesCollected) -> None: task_id = event.task_id if task_id not in self._pending_traces: self._pending_traces[task_id] = {} self._pending_traces[task_id][event.rank] = event.traces if ( task_id in self._expected_ranks and set(self._pending_traces[task_id].keys()) >= self._expected_ranks[task_id] ): await self._merge_and_save_traces(task_id) async def _merge_and_save_traces(self, task_id: TaskId) -> None: all_trace_data: list[TraceEventData] = [] for trace_data in self._pending_traces[task_id].values(): all_trace_data.extend(trace_data) await self.event_sender.send( TracesMerged(task_id=task_id, traces=all_trace_data) ) del self._pending_traces[task_id] if task_id in self._expected_ranks: del self._expected_ranks[task_id] ================================================ FILE: src/exo/master/placement.py ================================================ import random from collections.abc import Mapping from copy import deepcopy from typing import Sequence from exo.master.placement_utils import ( Cycle, filter_cycles_by_memory, get_mlx_jaccl_coordinators, get_mlx_jaccl_devices_matrix, get_mlx_ring_hosts_by_node, get_shard_assignments, get_smallest_cycles, ) from exo.shared.models.model_cards import ModelId from exo.shared.topology import Topology from exo.shared.types.commands import ( CancelDownload, CreateInstance, DeleteInstance, DownloadCommand, PlaceInstance, ) from exo.shared.types.common import NodeId from exo.shared.types.events import ( Event, InstanceCreated, InstanceDeleted, TaskStatusUpdated, ) from exo.shared.types.memory import Memory from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.worker.downloads import ( DownloadOngoing, DownloadProgress, ) from exo.shared.types.worker.instances import ( Instance, InstanceId, InstanceMeta, MlxJacclInstance, MlxRingInstance, ) from exo.shared.types.worker.shards import Sharding def random_ephemeral_port() -> int: port = random.randint(49153, 65535) return port - 1 if port <= 52415 else port def add_instance_to_placements( command: CreateInstance, topology: Topology, current_instances: Mapping[InstanceId, Instance], ) -> Mapping[InstanceId, Instance]: # TODO: validate against topology return {**current_instances, command.instance.instance_id: command.instance} def place_instance( command: PlaceInstance, topology: Topology, current_instances: Mapping[InstanceId, Instance], node_memory: Mapping[NodeId, MemoryUsage], node_network: Mapping[NodeId, NodeNetworkInfo], required_nodes: set[NodeId] | None = None, ) -> dict[InstanceId, Instance]: cycles = topology.get_cycles() candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles)) # Filter to cycles containing all required nodes (subset matching) if required_nodes: candidate_cycles = [ cycle for cycle in candidate_cycles if required_nodes.issubset(cycle.node_ids) ] cycles_with_sufficient_memory = filter_cycles_by_memory( candidate_cycles, node_memory, command.model_card.storage_size ) if len(cycles_with_sufficient_memory) == 0: raise ValueError("No cycles found with sufficient memory") if command.sharding == Sharding.Tensor: if not command.model_card.supports_tensor: raise ValueError( f"Requested Tensor sharding but this model does not support tensor parallelism: {command.model_card.model_id}" ) # TODO: the condition here for tensor parallel is not correct, but it works good enough for now. kv_heads = command.model_card.num_key_value_heads cycles_with_sufficient_memory = [ cycle for cycle in cycles_with_sufficient_memory if command.model_card.hidden_size % len(cycle) == 0 and (kv_heads is None or kv_heads % len(cycle) == 0) ] if not cycles_with_sufficient_memory: raise ValueError( f"No tensor sharding found for model with " f"hidden_size={command.model_card.hidden_size}" f"{f', num_key_value_heads={kv_heads}' if kv_heads is not None else ''}" f" across candidate cycles" ) if command.sharding == Sharding.Pipeline and command.model_card.model_id == ModelId( "mlx-community/DeepSeek-V3.1-8bit" ): raise ValueError( "Pipeline parallelism is not supported for DeepSeek V3.1 (8-bit)" ) smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory) smallest_rdma_cycles = [ cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle) ] if command.instance_meta == InstanceMeta.MlxJaccl: if not smallest_rdma_cycles: raise ValueError( "Requested RDMA (MlxJaccl) but no RDMA-connected cycles available" ) smallest_cycles = smallest_rdma_cycles cycles_with_leaf_nodes: list[Cycle] = [ cycle for cycle in smallest_cycles if any(topology.node_is_leaf(node_id) for node_id in cycle) ] selected_cycle = max( cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else smallest_cycles, key=lambda cycle: sum( (node_memory[node_id].ram_available for node_id in cycle), start=Memory(), ), ) # Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node) if len(selected_cycle) == 1: command.instance_meta = InstanceMeta.MlxRing command.sharding = Sharding.Pipeline shard_assignments = get_shard_assignments( command.model_card, selected_cycle, command.sharding, node_memory ) cycle_digraph: Topology = topology.get_subgraph_from_nodes(selected_cycle.node_ids) instance_id = InstanceId() target_instances = dict(deepcopy(current_instances)) match command.instance_meta: case InstanceMeta.MlxJaccl: # TODO(evan): shard assignments should contain information about ranks, this is ugly def get_device_rank(node_id: NodeId) -> int: runner_id = shard_assignments.node_to_runner[node_id] shard_metadata = shard_assignments.runner_to_shard.get(runner_id) assert shard_metadata is not None return shard_metadata.device_rank zero_node_ids = [ node_id for node_id in selected_cycle.node_ids if get_device_rank(node_id) == 0 ] assert len(zero_node_ids) == 1 coordinator_node_id = zero_node_ids[0] mlx_jaccl_devices = get_mlx_jaccl_devices_matrix( [node_id for node_id in selected_cycle], cycle_digraph, ) mlx_jaccl_coordinators = get_mlx_jaccl_coordinators( coordinator=coordinator_node_id, coordinator_port=random_ephemeral_port(), cycle_digraph=cycle_digraph, node_network=node_network, ) target_instances[instance_id] = MlxJacclInstance( instance_id=instance_id, shard_assignments=shard_assignments, jaccl_devices=mlx_jaccl_devices, jaccl_coordinators=mlx_jaccl_coordinators, ) case InstanceMeta.MlxRing: ephemeral_port = random_ephemeral_port() hosts_by_node = get_mlx_ring_hosts_by_node( selected_cycle=selected_cycle, cycle_digraph=cycle_digraph, ephemeral_port=ephemeral_port, node_network=node_network, ) target_instances[instance_id] = MlxRingInstance( instance_id=instance_id, shard_assignments=shard_assignments, hosts_by_node=hosts_by_node, ephemeral_port=ephemeral_port, ) return target_instances def delete_instance( command: DeleteInstance, current_instances: Mapping[InstanceId, Instance], ) -> dict[InstanceId, Instance]: target_instances = dict(deepcopy(current_instances)) if command.instance_id in target_instances: del target_instances[command.instance_id] return target_instances raise ValueError(f"Instance {command.instance_id} not found") def get_transition_events( current_instances: Mapping[InstanceId, Instance], target_instances: Mapping[InstanceId, Instance], tasks: Mapping[TaskId, Task], ) -> Sequence[Event]: events: list[Event] = [] # find instances to create for instance_id, instance in target_instances.items(): if instance_id not in current_instances: events.append( InstanceCreated( instance=instance, ) ) # find instances to delete for instance_id in current_instances: if instance_id not in target_instances: for task in tasks.values(): if task.instance_id == instance_id and task.task_status in [ TaskStatus.Pending, TaskStatus.Running, ]: events.append( TaskStatusUpdated( task_status=TaskStatus.Cancelled, task_id=task.task_id, ) ) events.append( InstanceDeleted( instance_id=instance_id, ) ) return events def cancel_unnecessary_downloads( instances: Mapping[InstanceId, Instance], download_status: Mapping[NodeId, Sequence[DownloadProgress]], ) -> Sequence[DownloadCommand]: commands: list[DownloadCommand] = [] currently_downloading = [ (k, v.shard_metadata.model_card.model_id) for k, vs in download_status.items() for v in vs if isinstance(v, (DownloadOngoing)) ] active_models = set( ( node_id, instance.shard_assignments.runner_to_shard[runner_id].model_card.model_id, ) for instance in instances.values() for node_id, runner_id in instance.shard_assignments.node_to_runner.items() ) for pair in currently_downloading: if pair not in active_models: commands.append(CancelDownload(target_node_id=pair[0], model_id=pair[1])) return commands ================================================ FILE: src/exo/master/placement_utils.py ================================================ from collections.abc import Generator, Mapping from loguru import logger from exo.shared.models.model_cards import ModelCard from exo.shared.topology import Topology from exo.shared.types.common import Host, NodeId from exo.shared.types.memory import Memory from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo from exo.shared.types.topology import Cycle, RDMAConnection, SocketConnection from exo.shared.types.worker.runners import RunnerId, ShardAssignments from exo.shared.types.worker.shards import ( CfgShardMetadata, PipelineShardMetadata, Sharding, ShardMetadata, TensorShardMetadata, ) def filter_cycles_by_memory( cycles: list[Cycle], node_memory: Mapping[NodeId, MemoryUsage], required_memory: Memory, ) -> list[Cycle]: filtered_cycles: list[Cycle] = [] for cycle in cycles: if not all(node in node_memory for node in cycle): continue total_mem = sum( (node_memory[node_id].ram_available for node_id in cycle.node_ids), start=Memory(), ) if total_mem >= required_memory: filtered_cycles.append(cycle) return filtered_cycles def get_smallest_cycles( cycles: list[Cycle], ) -> list[Cycle]: min_nodes = min(len(cycle) for cycle in cycles) return [cycle for cycle in cycles if len(cycle) == min_nodes] def allocate_layers_proportionally( total_layers: int, memory_fractions: list[float], ) -> list[int]: n = len(memory_fractions) if n == 0: raise ValueError("Cannot allocate layers to an empty node list") if total_layers < n: raise ValueError( f"Cannot distribute {total_layers} layers across {n} nodes " "(need at least 1 layer per node)" ) # Largest remainder: floor each, then distribute remainder by fractional part raw = [f * total_layers for f in memory_fractions] result = [int(r) for r in raw] by_remainder = sorted(range(n), key=lambda i: raw[i] - result[i], reverse=True) for i in range(total_layers - sum(result)): result[by_remainder[i]] += 1 # Ensure minimum 1 per node by taking from the largest for i in range(n): if result[i] == 0: max_idx = max(range(n), key=lambda j: result[j]) assert result[max_idx] > 1 result[max_idx] -= 1 result[i] = 1 return result def _validate_cycle(cycle: Cycle) -> None: if not cycle.node_ids: raise ValueError("Cannot create shard assignments for empty node cycle") def _compute_total_memory( node_ids: list[NodeId], node_memory: Mapping[NodeId, MemoryUsage], ) -> Memory: total_memory = sum( (node_memory[node_id].ram_available for node_id in node_ids), start=Memory(), ) if total_memory.in_bytes == 0: raise ValueError("Cannot create shard assignments: total available memory is 0") return total_memory def _allocate_and_validate_layers( node_ids: list[NodeId], node_memory: Mapping[NodeId, MemoryUsage], total_memory: Memory, model_card: ModelCard, ) -> list[int]: layer_allocations = allocate_layers_proportionally( total_layers=model_card.n_layers, memory_fractions=[ node_memory[node_id].ram_available / total_memory for node_id in node_ids ], ) total_storage = model_card.storage_size total_layers = model_card.n_layers for i, node_id in enumerate(node_ids): node_layers = layer_allocations[i] required_memory = (total_storage * node_layers) // total_layers available_memory = node_memory[node_id].ram_available if required_memory > available_memory: raise ValueError( f"Node {i} ({node_id}) has insufficient memory: " f"requires {required_memory.in_gb:.2f} GB for {node_layers} layers, " f"but only has {available_memory.in_gb:.2f} GB available" ) return layer_allocations def get_shard_assignments_for_pipeline_parallel( model_card: ModelCard, cycle: Cycle, node_memory: Mapping[NodeId, MemoryUsage], ) -> ShardAssignments: """Create shard assignments for pipeline parallel execution.""" world_size = len(cycle) use_cfg_parallel = model_card.uses_cfg and world_size >= 2 and world_size % 2 == 0 if use_cfg_parallel: return _get_shard_assignments_for_cfg_parallel(model_card, cycle, node_memory) else: return _get_shard_assignments_for_pure_pipeline(model_card, cycle, node_memory) def _get_shard_assignments_for_cfg_parallel( model_card: ModelCard, cycle: Cycle, node_memory: Mapping[NodeId, MemoryUsage], ) -> ShardAssignments: """Create shard assignments for CFG parallel execution. CFG parallel runs two independent pipelines. Group 0 processes the positive prompt, group 1 processes the negative prompt. The ring topology places group 1's ranks in reverse order so both "last stages" are neighbors for efficient CFG exchange. """ _validate_cycle(cycle) world_size = len(cycle) cfg_world_size = 2 pipeline_world_size = world_size // cfg_world_size # Allocate layers for one pipeline group (both groups run the same layers) pipeline_node_ids = cycle.node_ids[:pipeline_world_size] pipeline_memory = _compute_total_memory(pipeline_node_ids, node_memory) layer_allocations = _allocate_and_validate_layers( pipeline_node_ids, node_memory, pipeline_memory, model_card ) # Ring topology: group 0 ascending [0,1,2,...], group 1 descending [...,2,1,0] # This places both last stages as neighbors for CFG exchange. position_to_cfg_pipeline = [(0, r) for r in range(pipeline_world_size)] + [ (1, r) for r in reversed(range(pipeline_world_size)) ] runner_to_shard: dict[RunnerId, ShardMetadata] = {} node_to_runner: dict[NodeId, RunnerId] = {} for device_rank, node_id in enumerate(cycle.node_ids): cfg_rank, pipeline_rank = position_to_cfg_pipeline[device_rank] layers_before = sum(layer_allocations[:pipeline_rank]) node_layers = layer_allocations[pipeline_rank] shard = CfgShardMetadata( model_card=model_card, device_rank=device_rank, world_size=world_size, start_layer=layers_before, end_layer=layers_before + node_layers, n_layers=model_card.n_layers, cfg_rank=cfg_rank, cfg_world_size=cfg_world_size, pipeline_rank=pipeline_rank, pipeline_world_size=pipeline_world_size, ) runner_id = RunnerId() runner_to_shard[runner_id] = shard node_to_runner[node_id] = runner_id return ShardAssignments( model_id=model_card.model_id, runner_to_shard=runner_to_shard, node_to_runner=node_to_runner, ) def _get_shard_assignments_for_pure_pipeline( model_card: ModelCard, cycle: Cycle, node_memory: Mapping[NodeId, MemoryUsage], ) -> ShardAssignments: """Create shard assignments for pure pipeline execution.""" _validate_cycle(cycle) total_memory = _compute_total_memory(cycle.node_ids, node_memory) layer_allocations = _allocate_and_validate_layers( cycle.node_ids, node_memory, total_memory, model_card ) runner_to_shard: dict[RunnerId, ShardMetadata] = {} node_to_runner: dict[NodeId, RunnerId] = {} for pipeline_rank, node_id in enumerate(cycle.node_ids): layers_before = sum(layer_allocations[:pipeline_rank]) node_layers = layer_allocations[pipeline_rank] shard = PipelineShardMetadata( model_card=model_card, device_rank=pipeline_rank, world_size=len(cycle), start_layer=layers_before, end_layer=layers_before + node_layers, n_layers=model_card.n_layers, ) runner_id = RunnerId() runner_to_shard[runner_id] = shard node_to_runner[node_id] = runner_id return ShardAssignments( model_id=model_card.model_id, runner_to_shard=runner_to_shard, node_to_runner=node_to_runner, ) def get_shard_assignments_for_tensor_parallel( model_card: ModelCard, cycle: Cycle, ): total_layers = model_card.n_layers world_size = len(cycle) runner_to_shard: dict[RunnerId, ShardMetadata] = {} node_to_runner: dict[NodeId, RunnerId] = {} for i, node_id in enumerate(cycle): shard = TensorShardMetadata( model_card=model_card, device_rank=i, world_size=world_size, start_layer=0, end_layer=total_layers, n_layers=total_layers, ) runner_id = RunnerId() runner_to_shard[runner_id] = shard node_to_runner[node_id] = runner_id shard_assignments = ShardAssignments( model_id=model_card.model_id, runner_to_shard=runner_to_shard, node_to_runner=node_to_runner, ) return shard_assignments def get_shard_assignments( model_card: ModelCard, cycle: Cycle, sharding: Sharding, node_memory: Mapping[NodeId, MemoryUsage], ) -> ShardAssignments: match sharding: case Sharding.Pipeline: return get_shard_assignments_for_pipeline_parallel( model_card=model_card, cycle=cycle, node_memory=node_memory, ) case Sharding.Tensor: return get_shard_assignments_for_tensor_parallel( model_card=model_card, cycle=cycle, ) def get_mlx_jaccl_devices_matrix( selected_cycle: list[NodeId], cycle_digraph: Topology, ) -> list[list[str | None]]: """Build connectivity matrix mapping device i to device j via RDMA interface names. The matrix element [i][j] contains the interface name on device i that connects to device j, or None if no connection exists or no interface name is found. Diagonal elements are always None. """ num_nodes = len(selected_cycle) matrix: list[list[str | None]] = [ [None for _ in range(num_nodes)] for _ in range(num_nodes) ] for i, node_i in enumerate(selected_cycle): for j, node_j in enumerate(selected_cycle): if i == j: continue for conn in cycle_digraph.get_all_connections_between(node_i, node_j): if isinstance(conn, RDMAConnection): matrix[i][j] = conn.source_rdma_iface break else: raise ValueError( "Current jaccl backend requires all-to-all RDMA connections" ) return matrix def _find_connection_ip( node_i: NodeId, node_j: NodeId, cycle_digraph: Topology, ) -> Generator[str, None, None]: """Find all IP addresses that connect node i to node j.""" for connection in cycle_digraph.get_all_connections_between(node_i, node_j): if isinstance(connection, SocketConnection): yield connection.sink_multiaddr.ip_address def _find_ip_prioritised( node_id: NodeId, other_node_id: NodeId, cycle_digraph: Topology, node_network: Mapping[NodeId, NodeNetworkInfo], ring: bool, ) -> str | None: """Find an IP address between nodes with prioritization. Priority: ethernet > wifi > unknown > thunderbolt """ ips = list(_find_connection_ip(node_id, other_node_id, cycle_digraph)) if not ips: return None other_network = node_network.get(other_node_id, NodeNetworkInfo()) ip_to_type = { iface.ip_address: iface.interface_type for iface in other_network.interfaces } # Ring should prioritise fastest connection. As a best-effort, we prioritise TB. # TODO: Profile and get actual connection speeds. if ring: priority = { "thunderbolt": 0, "maybe_ethernet": 1, "ethernet": 2, "wifi": 3, "unknown": 4, } # RDMA prefers ethernet coordinator else: priority = { "ethernet": 0, "wifi": 1, "unknown": 2, "maybe_ethernet": 3, "thunderbolt": 4, } return min(ips, key=lambda ip: priority.get(ip_to_type.get(ip, "unknown"), 2)) def get_mlx_ring_hosts_by_node( selected_cycle: Cycle, cycle_digraph: Topology, ephemeral_port: int, node_network: Mapping[NodeId, NodeNetworkInfo], ) -> dict[NodeId, list[Host]]: """Generate per-node host lists for MLX ring backend. Each node gets a list where: - Self position: Host(ip="0.0.0.0", port=ephemeral_port) - Left/right neighbors: actual connection IPs - Non-neighbors: Host(ip="198.51.100.1", port=0) placeholder (RFC 5737 TEST-NET-2) """ world_size = len(selected_cycle) if world_size == 0: return {} hosts_by_node: dict[NodeId, list[Host]] = {} for rank, node_id in enumerate(selected_cycle): left_rank = (rank - 1) % world_size right_rank = (rank + 1) % world_size hosts_for_node: list[Host] = [] for idx, other_node_id in enumerate(selected_cycle): if idx == rank: hosts_for_node.append(Host(ip="0.0.0.0", port=ephemeral_port)) continue if idx not in {left_rank, right_rank}: # Placeholder IP from RFC 5737 TEST-NET-2 hosts_for_node.append(Host(ip="198.51.100.1", port=0)) continue connection_ip = _find_ip_prioritised( node_id, other_node_id, cycle_digraph, node_network, ring=True ) if connection_ip is None: raise ValueError( "MLX ring backend requires connectivity between neighbouring nodes" ) hosts_for_node.append(Host(ip=connection_ip, port=ephemeral_port)) hosts_by_node[node_id] = hosts_for_node return hosts_by_node def get_mlx_jaccl_coordinators( coordinator: NodeId, coordinator_port: int, cycle_digraph: Topology, node_network: Mapping[NodeId, NodeNetworkInfo], ) -> dict[NodeId, str]: """Get the coordinator addresses for MLX JACCL (rank 0 device). Select an IP address that each node can reach for the rank 0 node. Returns address in format "X.X.X.X:PORT" per node. """ logger.debug(f"Selecting coordinator: {coordinator}") def get_ip_for_node(n: NodeId) -> str: if n == coordinator: return "0.0.0.0" ip = _find_ip_prioritised( n, coordinator, cycle_digraph, node_network, ring=False ) if ip is not None: return ip raise ValueError( "Current jaccl backend requires all participating devices to be able to communicate" ) return { n: f"{get_ip_for_node(n)}:{coordinator_port}" for n in cycle_digraph.list_nodes() } ================================================ FILE: src/exo/master/tests/conftest.py ================================================ from exo.shared.types.multiaddr import Multiaddr from exo.shared.types.profiling import ( MemoryUsage, NetworkInterfaceInfo, NodeNetworkInfo, ) from exo.shared.types.topology import RDMAConnection, SocketConnection def create_node_memory(memory: int) -> MemoryUsage: return MemoryUsage.from_bytes( ram_total=1000, ram_available=memory, swap_total=1000, swap_available=1000, ) def create_node_network() -> NodeNetworkInfo: return NodeNetworkInfo( interfaces=[ NetworkInterfaceInfo(name="en0", ip_address=f"169.254.0.{i}") for i in range(10) ] ) def create_socket_connection(ip: int, sink_port: int = 1234) -> SocketConnection: return SocketConnection( sink_multiaddr=Multiaddr(address=f"/ip4/169.254.0.{ip}/tcp/{sink_port}"), ) def create_rdma_connection(iface: int) -> RDMAConnection: return RDMAConnection( source_rdma_iface=f"rdma_en{iface}", sink_rdma_iface=f"rdma_en{iface}" ) ================================================ FILE: src/exo/master/tests/test_master.py ================================================ from datetime import datetime, timezone from typing import Sequence import anyio import pytest from loguru import logger from exo.master.main import Master from exo.routing.router import get_node_id_keypair from exo.shared.models.model_cards import ModelCard, ModelTask from exo.shared.types.commands import ( CommandId, ForwarderCommand, ForwarderDownloadCommand, PlaceInstance, TextGeneration, ) from exo.shared.types.common import ModelId, NodeId, SessionId, SystemId from exo.shared.types.events import ( Event, GlobalForwarderEvent, IndexedEvent, InstanceCreated, LocalForwarderEvent, NodeGatheredInfo, TaskCreated, ) from exo.shared.types.memory import Memory from exo.shared.types.profiling import ( MemoryUsage, ) from exo.shared.types.tasks import TaskStatus from exo.shared.types.tasks import TextGeneration as TextGenerationTask from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams from exo.shared.types.worker.instances import ( InstanceMeta, MlxRingInstance, ShardAssignments, ) from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding from exo.utils.channels import channel @pytest.mark.asyncio async def test_master(): keypair = get_node_id_keypair() node_id = NodeId(keypair.to_node_id()) session_id = SessionId(master_node_id=node_id, election_clock=0) ge_sender, global_event_receiver = channel[GlobalForwarderEvent]() command_sender, co_receiver = channel[ForwarderCommand]() local_event_sender, le_receiver = channel[LocalForwarderEvent]() fcds, _fcdr = channel[ForwarderDownloadCommand]() ev_send, ev_recv = channel[Event]() async def mock_event_router(): idx = 0 sid = SystemId() with ev_recv as master_events: async for event in master_events: await local_event_sender.send( LocalForwarderEvent( origin=sid, origin_idx=idx, session=session_id, event=event, ) ) idx += 1 all_events: list[IndexedEvent] = [] def _get_events() -> Sequence[IndexedEvent]: orig_events = global_event_receiver.collect() for e in orig_events: all_events.append( IndexedEvent( event=e.event, idx=len(all_events), # origin=e.origin, ) ) return all_events master = Master( node_id, session_id, event_sender=ev_send, global_event_sender=ge_sender, local_event_receiver=le_receiver, command_receiver=co_receiver, download_command_sender=fcds, ) logger.info("run the master") async with anyio.create_task_group() as tg: tg.start_soon(master.run) tg.start_soon(mock_event_router) # inject a NodeGatheredInfo event logger.info("inject a NodeGatheredInfo event") await local_event_sender.send( LocalForwarderEvent( origin_idx=0, origin=SystemId("Worker"), session=session_id, event=( NodeGatheredInfo( when=str(datetime.now(tz=timezone.utc)), node_id=node_id, info=MemoryUsage( ram_total=Memory.from_bytes(678948 * 1024), ram_available=Memory.from_bytes(678948 * 1024), swap_total=Memory.from_bytes(0), swap_available=Memory.from_bytes(0), ), ) ), ) ) # wait for initial topology event logger.info("wait for initial topology event") while len(list(master.state.topology.list_nodes())) == 0: await anyio.sleep(0.001) while len(master.state.node_memory) == 0: await anyio.sleep(0.001) logger.info("inject a CreateInstance Command") await command_sender.send( ForwarderCommand( origin=SystemId("API"), command=( PlaceInstance( command_id=CommandId(), model_card=ModelCard( model_id=ModelId("llama-3.2-1b"), n_layers=16, storage_size=Memory.from_bytes(678948), hidden_size=7168, supports_tensor=True, tasks=[ModelTask.TextGeneration], ), sharding=Sharding.Pipeline, instance_meta=InstanceMeta.MlxRing, min_nodes=1, ) ), ) ) logger.info("wait for an instance") while len(master.state.instances.keys()) == 0: await anyio.sleep(0.001) logger.info("inject a TextGeneration Command") await command_sender.send( ForwarderCommand( origin=SystemId("API"), command=( TextGeneration( command_id=CommandId(), task_params=TextGenerationTaskParams( model=ModelId("llama-3.2-1b"), input=[ InputMessage(role="user", content="Hello, how are you?") ], ), ) ), ) ) while len(_get_events()) < 3: await anyio.sleep(0.01) events = _get_events() assert len(events) == 3 assert events[0].idx == 0 assert events[1].idx == 1 assert events[2].idx == 2 assert isinstance(events[0].event, NodeGatheredInfo) assert isinstance(events[1].event, InstanceCreated) created_instance = events[1].event.instance assert isinstance(created_instance, MlxRingInstance) runner_id = list(created_instance.shard_assignments.runner_to_shard.keys())[0] # Validate the shard assignments expected_shard_assignments = ShardAssignments( model_id=ModelId("llama-3.2-1b"), runner_to_shard={ (runner_id): PipelineShardMetadata( start_layer=0, end_layer=16, n_layers=16, model_card=ModelCard( model_id=ModelId("llama-3.2-1b"), n_layers=16, storage_size=Memory.from_bytes(678948), hidden_size=7168, supports_tensor=True, tasks=[ModelTask.TextGeneration], ), device_rank=0, world_size=1, ) }, node_to_runner={node_id: runner_id}, ) assert created_instance.shard_assignments == expected_shard_assignments # For single-node, hosts_by_node should have one entry with self-binding assert len(created_instance.hosts_by_node) == 1 assert node_id in created_instance.hosts_by_node assert len(created_instance.hosts_by_node[node_id]) == 1 assert created_instance.hosts_by_node[node_id][0].ip == "0.0.0.0" assert created_instance.ephemeral_port > 0 assert isinstance(events[2].event, TaskCreated) assert events[2].event.task.task_status == TaskStatus.Pending assert isinstance(events[2].event.task, TextGenerationTask) assert events[2].event.task.task_params == TextGenerationTaskParams( model=ModelId("llama-3.2-1b"), input=[InputMessage(role="user", content="Hello, how are you?")], ) ev_send.close() await master.shutdown() ================================================ FILE: src/exo/master/tests/test_placement.py ================================================ import pytest from exo.master.placement import ( get_transition_events, place_instance, ) from exo.master.tests.conftest import ( create_node_memory, create_node_network, create_rdma_connection, create_socket_connection, ) from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask from exo.shared.topology import Topology from exo.shared.types.commands import PlaceInstance from exo.shared.types.common import CommandId, NodeId from exo.shared.types.events import ( InstanceCreated, InstanceDeleted, TaskStatusUpdated, ) from exo.shared.types.memory import Memory from exo.shared.types.multiaddr import Multiaddr from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams from exo.shared.types.topology import Connection, SocketConnection from exo.shared.types.worker.instances import ( Instance, InstanceId, InstanceMeta, MlxJacclInstance, MlxRingInstance, ) from exo.shared.types.worker.runners import ShardAssignments from exo.shared.types.worker.shards import Sharding @pytest.fixture def instance() -> Instance: return MlxRingInstance( instance_id=InstanceId(), shard_assignments=ShardAssignments( model_id=ModelId("test-model"), runner_to_shard={}, node_to_runner={} ), hosts_by_node={}, ephemeral_port=50000, ) @pytest.fixture def model_card() -> ModelCard: return ModelCard( model_id=ModelId("test-model"), storage_size=Memory.from_kb(1000), n_layers=10, hidden_size=30, supports_tensor=True, tasks=[ModelTask.TextGeneration], ) def place_instance_command(model_card: ModelCard) -> PlaceInstance: return PlaceInstance( command_id=CommandId(), model_card=model_card, sharding=Sharding.Pipeline, instance_meta=InstanceMeta.MlxRing, min_nodes=1, ) @pytest.mark.parametrize( "available_memory,total_layers,expected_layers", [ ((500, 500, 1000), 12, (3, 3, 6)), ((500, 500, 500), 12, (4, 4, 4)), ((312, 468, 1092), 12, (2, 3, 7)), ], ) def test_get_instance_placements_create_instance( available_memory: tuple[int, int, int], total_layers: int, expected_layers: tuple[int, int, int], model_card: ModelCard, ): # arrange model_card.n_layers = total_layers model_card.storage_size = Memory.from_bytes( sum(available_memory) ) # make it exactly fit across all nodes topology = Topology() cic = place_instance_command(model_card) node_id_a = NodeId() node_id_b = NodeId() node_id_c = NodeId() # fully connected (directed) between the 3 nodes conn_a_b = Connection( source=node_id_a, sink=node_id_b, edge=create_socket_connection(1) ) conn_b_c = Connection( source=node_id_b, sink=node_id_c, edge=create_socket_connection(2) ) conn_c_a = Connection( source=node_id_c, sink=node_id_a, edge=create_socket_connection(3) ) conn_c_b = Connection( source=node_id_c, sink=node_id_b, edge=create_socket_connection(4) ) conn_a_c = Connection( source=node_id_a, sink=node_id_c, edge=create_socket_connection(5) ) conn_b_a = Connection( source=node_id_b, sink=node_id_a, edge=create_socket_connection(6) ) node_memory = { node_id_a: create_node_memory(available_memory[0]), node_id_b: create_node_memory(available_memory[1]), node_id_c: create_node_memory(available_memory[2]), } node_network = { node_id_a: create_node_network(), node_id_b: create_node_network(), node_id_c: create_node_network(), } topology.add_node(node_id_a) topology.add_node(node_id_b) topology.add_node(node_id_c) topology.add_connection(conn_a_b) topology.add_connection(conn_b_c) topology.add_connection(conn_c_a) topology.add_connection(conn_c_b) topology.add_connection(conn_a_c) topology.add_connection(conn_b_a) # act placements = place_instance(cic, topology, {}, node_memory, node_network) # assert assert len(placements) == 1 instance_id = list(placements.keys())[0] instance = placements[instance_id] assert instance.shard_assignments.model_id == model_card.model_id runner_id_a = instance.shard_assignments.node_to_runner[node_id_a] runner_id_b = instance.shard_assignments.node_to_runner[node_id_b] runner_id_c = instance.shard_assignments.node_to_runner[node_id_c] shard_a = instance.shard_assignments.runner_to_shard[runner_id_a] shard_b = instance.shard_assignments.runner_to_shard[runner_id_b] shard_c = instance.shard_assignments.runner_to_shard[runner_id_c] assert shard_a.end_layer - shard_a.start_layer == expected_layers[0] assert shard_b.end_layer - shard_b.start_layer == expected_layers[1] assert shard_c.end_layer - shard_c.start_layer == expected_layers[2] shards = [shard_a, shard_b, shard_c] shards_sorted = sorted(shards, key=lambda s: s.start_layer) assert shards_sorted[0].start_layer == 0 assert shards_sorted[-1].end_layer == total_layers def test_get_instance_placements_one_node_exact_fit() -> None: topology = Topology() node_id = NodeId() topology.add_node(node_id) node_memory = {node_id: create_node_memory(1000 * 1024)} node_network = {node_id: create_node_network()} cic = place_instance_command( ModelCard( model_id=ModelId("test-model"), storage_size=Memory.from_kb(1000), n_layers=10, hidden_size=1000, supports_tensor=True, tasks=[ModelTask.TextGeneration], ), ) placements = place_instance(cic, topology, {}, node_memory, node_network) assert len(placements) == 1 instance_id = list(placements.keys())[0] instance = placements[instance_id] assert instance.shard_assignments.model_id == "test-model" assert len(instance.shard_assignments.node_to_runner) == 1 assert len(instance.shard_assignments.runner_to_shard) == 1 assert len(instance.shard_assignments.runner_to_shard) == 1 def test_get_instance_placements_one_node_fits_with_extra_memory() -> None: topology = Topology() node_id = NodeId() topology.add_node(node_id) node_memory = {node_id: create_node_memory(1001 * 1024)} node_network = {node_id: create_node_network()} cic = place_instance_command( ModelCard( model_id=ModelId("test-model"), storage_size=Memory.from_kb(1000), n_layers=10, hidden_size=1000, supports_tensor=True, tasks=[ModelTask.TextGeneration], ), ) placements = place_instance(cic, topology, {}, node_memory, node_network) assert len(placements) == 1 instance_id = list(placements.keys())[0] instance = placements[instance_id] assert instance.shard_assignments.model_id == "test-model" assert len(instance.shard_assignments.node_to_runner) == 1 assert len(instance.shard_assignments.runner_to_shard) == 1 assert len(instance.shard_assignments.runner_to_shard) == 1 def test_get_instance_placements_one_node_not_fit() -> None: topology = Topology() node_id = NodeId() topology.add_node(node_id) node_memory = {node_id: create_node_memory(1000 * 1024)} node_network = {node_id: create_node_network()} cic = place_instance_command( model_card=ModelCard( model_id=ModelId("test-model"), storage_size=Memory.from_kb(1001), n_layers=10, hidden_size=1000, supports_tensor=True, tasks=[ModelTask.TextGeneration], ), ) with pytest.raises(ValueError, match="No cycles found with sufficient memory"): place_instance(cic, topology, {}, node_memory, node_network) def test_get_transition_events_no_change(instance: Instance): # arrange instance_id = InstanceId() current_instances = {instance_id: instance} target_instances = {instance_id: instance} # act events = get_transition_events(current_instances, target_instances, {}) # assert assert len(events) == 0 def test_get_transition_events_create_instance(instance: Instance): # arrange instance_id = InstanceId() current_instances: dict[InstanceId, Instance] = {} target_instances: dict[InstanceId, Instance] = {instance_id: instance} # act events = get_transition_events(current_instances, target_instances, {}) # assert assert len(events) == 1 assert isinstance(events[0], InstanceCreated) def test_get_transition_events_delete_instance(instance: Instance): # arrange instance_id = InstanceId() current_instances: dict[InstanceId, Instance] = {instance_id: instance} target_instances: dict[InstanceId, Instance] = {} # act events = get_transition_events(current_instances, target_instances, {}) # assert assert len(events) == 1 assert isinstance(events[0], InstanceDeleted) assert events[0].instance_id == instance_id def test_placement_selects_leaf_nodes( model_card: ModelCard, ): # arrange topology = Topology() model_card.storage_size = Memory.from_bytes(1000) node_id_a = NodeId() node_id_b = NodeId() node_id_c = NodeId() node_id_d = NodeId() node_memory = { node_id_a: create_node_memory(500), node_id_b: create_node_memory(600), node_id_c: create_node_memory(600), node_id_d: create_node_memory(500), } node_network = { node_id_a: create_node_network(), node_id_b: create_node_network(), node_id_c: create_node_network(), node_id_d: create_node_network(), } topology.add_node(node_id_a) topology.add_node(node_id_b) topology.add_node(node_id_c) topology.add_node(node_id_d) # Daisy chain topology (directed) topology.add_connection( Connection(source=node_id_a, sink=node_id_b, edge=create_socket_connection(1)) ) topology.add_connection( Connection(source=node_id_b, sink=node_id_a, edge=create_socket_connection(1)) ) topology.add_connection( Connection(source=node_id_b, sink=node_id_c, edge=create_socket_connection(1)) ) topology.add_connection( Connection(source=node_id_c, sink=node_id_b, edge=create_socket_connection(1)) ) topology.add_connection( Connection(source=node_id_c, sink=node_id_d, edge=create_socket_connection(1)) ) topology.add_connection( Connection(source=node_id_d, sink=node_id_c, edge=create_socket_connection(1)) ) cic = place_instance_command(model_card=model_card) # act placements = place_instance(cic, topology, {}, node_memory, node_network) # assert assert len(placements) == 1 instance = list(placements.values())[0] assigned_nodes = set(instance.shard_assignments.node_to_runner.keys()) assert assigned_nodes == set((node_id_a, node_id_b)) or assigned_nodes == set( ( node_id_c, node_id_d, ) ) def test_tensor_rdma_backend_connectivity_matrix( model_card: ModelCard, ): # arrange topology = Topology() model_card.n_layers = 12 model_card.storage_size = Memory.from_bytes(1500) node_a = NodeId() node_b = NodeId() node_c = NodeId() node_memory = { node_a: create_node_memory(500), node_b: create_node_memory(500), node_c: create_node_memory(500), } ethernet_interface = NetworkInterfaceInfo( name="en0", ip_address="10.0.0.1", ) ethernet_conn = SocketConnection( sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000") ) node_network = { node_a: NodeNetworkInfo(interfaces=[ethernet_interface]), node_b: NodeNetworkInfo(interfaces=[ethernet_interface]), node_c: NodeNetworkInfo(interfaces=[ethernet_interface]), } topology.add_node(node_a) topology.add_node(node_b) topology.add_node(node_c) # RDMA connections (directed) topology.add_connection( Connection(source=node_a, sink=node_b, edge=create_rdma_connection(3)) ) topology.add_connection( Connection(source=node_b, sink=node_a, edge=create_rdma_connection(3)) ) topology.add_connection( Connection(source=node_b, sink=node_c, edge=create_rdma_connection(4)) ) topology.add_connection( Connection(source=node_c, sink=node_b, edge=create_rdma_connection(4)) ) topology.add_connection( Connection(source=node_a, sink=node_c, edge=create_rdma_connection(5)) ) topology.add_connection( Connection(source=node_c, sink=node_a, edge=create_rdma_connection(5)) ) # Ethernet connections (directed) topology.add_connection(Connection(source=node_a, sink=node_b, edge=ethernet_conn)) topology.add_connection(Connection(source=node_b, sink=node_c, edge=ethernet_conn)) topology.add_connection(Connection(source=node_c, sink=node_a, edge=ethernet_conn)) topology.add_connection(Connection(source=node_a, sink=node_c, edge=ethernet_conn)) topology.add_connection(Connection(source=node_b, sink=node_a, edge=ethernet_conn)) topology.add_connection(Connection(source=node_c, sink=node_b, edge=ethernet_conn)) cic = PlaceInstance( sharding=Sharding.Tensor, instance_meta=InstanceMeta.MlxJaccl, command_id=CommandId(), model_card=model_card, min_nodes=1, ) # act placements = place_instance(cic, topology, {}, node_memory, node_network) # assert assert len(placements) == 1 instance_id = list(placements.keys())[0] instance = placements[instance_id] assert isinstance(instance, MlxJacclInstance) assert instance.jaccl_devices is not None assert instance.jaccl_coordinators is not None matrix = instance.jaccl_devices assert len(matrix) == 3 for i in range(3): assert matrix[i][i] is None assigned_nodes = list(instance.shard_assignments.node_to_runner.keys()) node_to_idx = {node_id: idx for idx, node_id in enumerate(assigned_nodes)} idx_a = node_to_idx[node_a] idx_b = node_to_idx[node_b] idx_c = node_to_idx[node_c] assert matrix[idx_a][idx_b] == "rdma_en3" assert matrix[idx_b][idx_c] == "rdma_en4" assert matrix[idx_c][idx_a] == "rdma_en5" # Verify coordinators are set for all nodes assert len(instance.jaccl_coordinators) == 3 for node_id in assigned_nodes: assert node_id in instance.jaccl_coordinators coordinator = instance.jaccl_coordinators[node_id] assert ":" in coordinator # Rank 0 node should use 0.0.0.0, others should use connection-specific IPs if node_id == assigned_nodes[0]: assert coordinator.startswith("0.0.0.0:") else: ip_part = coordinator.split(":")[0] assert len(ip_part.split(".")) == 4 def _make_task( instance_id: InstanceId, status: TaskStatus = TaskStatus.Running, ) -> TextGeneration: return TextGeneration( task_id=TaskId(), task_status=status, instance_id=instance_id, command_id=CommandId(), task_params=TextGenerationTaskParams( model=ModelId("test-model"), input=[InputMessage(role="user", content="hello")], ), ) def test_get_transition_events_delete_instance_cancels_running_tasks( instance: Instance, ): # arrange instance_id = InstanceId() current_instances: dict[InstanceId, Instance] = {instance_id: instance} target_instances: dict[InstanceId, Instance] = {} task = _make_task(instance_id, TaskStatus.Running) tasks = {task.task_id: task} # act events = get_transition_events(current_instances, target_instances, tasks) # assert – cancellation event should come before the deletion event assert len(events) == 2 assert isinstance(events[0], TaskStatusUpdated) assert events[0].task_id == task.task_id assert events[0].task_status == TaskStatus.Cancelled assert isinstance(events[1], InstanceDeleted) assert events[1].instance_id == instance_id def test_get_transition_events_delete_instance_cancels_pending_tasks( instance: Instance, ): # arrange instance_id = InstanceId() current_instances: dict[InstanceId, Instance] = {instance_id: instance} target_instances: dict[InstanceId, Instance] = {} task = _make_task(instance_id, TaskStatus.Pending) tasks = {task.task_id: task} # act events = get_transition_events(current_instances, target_instances, tasks) # assert assert len(events) == 2 assert isinstance(events[0], TaskStatusUpdated) assert events[0].task_id == task.task_id assert events[0].task_status == TaskStatus.Cancelled assert isinstance(events[1], InstanceDeleted) def test_get_transition_events_delete_instance_ignores_completed_tasks( instance: Instance, ): # arrange instance_id = InstanceId() current_instances: dict[InstanceId, Instance] = {instance_id: instance} target_instances: dict[InstanceId, Instance] = {} tasks = { t.task_id: t for t in [ _make_task(instance_id, TaskStatus.Complete), _make_task(instance_id, TaskStatus.Failed), _make_task(instance_id, TaskStatus.TimedOut), _make_task(instance_id, TaskStatus.Cancelled), ] } # act events = get_transition_events(current_instances, target_instances, tasks) # assert – only the InstanceDeleted event, no cancellations assert len(events) == 1 assert isinstance(events[0], InstanceDeleted) def test_get_transition_events_delete_instance_cancels_only_matching_tasks( instance: Instance, ): # arrange instance_id_a = InstanceId() instance_id_b = InstanceId() current_instances: dict[InstanceId, Instance] = { instance_id_a: instance, instance_id_b: instance, } # only delete instance A, keep instance B target_instances: dict[InstanceId, Instance] = {instance_id_b: instance} task_a = _make_task(instance_id_a, TaskStatus.Running) task_b = _make_task(instance_id_b, TaskStatus.Running) tasks = {task_a.task_id: task_a, task_b.task_id: task_b} # act events = get_transition_events(current_instances, target_instances, tasks) # assert – only task_a should be cancelled cancel_events = [e for e in events if isinstance(e, TaskStatusUpdated)] delete_events = [e for e in events if isinstance(e, InstanceDeleted)] assert len(cancel_events) == 1 assert cancel_events[0].task_id == task_a.task_id assert cancel_events[0].task_status == TaskStatus.Cancelled assert len(delete_events) == 1 assert delete_events[0].instance_id == instance_id_a ================================================ FILE: src/exo/master/tests/test_placement_utils.py ================================================ import pytest from exo.master.placement_utils import ( allocate_layers_proportionally, filter_cycles_by_memory, get_mlx_jaccl_coordinators, get_shard_assignments, get_shard_assignments_for_pipeline_parallel, get_smallest_cycles, ) from exo.master.tests.conftest import ( create_node_memory, create_socket_connection, ) from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask from exo.shared.topology import Topology from exo.shared.types.common import NodeId from exo.shared.types.memory import Memory from exo.shared.types.profiling import ( NetworkInterfaceInfo, NodeNetworkInfo, ) from exo.shared.types.topology import Connection, SocketConnection from exo.shared.types.worker.shards import ( CfgShardMetadata, PipelineShardMetadata, Sharding, ) def test_filter_cycles_by_memory(): # arrange node1_id = NodeId() node2_id = NodeId() connection1 = Connection( source=node1_id, sink=node2_id, edge=create_socket_connection(1) ) connection2 = Connection( source=node2_id, sink=node1_id, edge=create_socket_connection(2) ) node1_mem = create_node_memory(1000 * 1024) node2_mem = create_node_memory(1000 * 1024) node_memory = {node1_id: node1_mem, node2_id: node2_mem} topology = Topology() topology.add_node(node1_id) topology.add_node(node2_id) topology.add_connection(connection1) topology.add_connection(connection2) cycles = [c for c in topology.get_cycles() if len(c) != 1] assert len(cycles) == 1 assert len(cycles[0]) == 2 # act filtered_cycles = filter_cycles_by_memory(cycles, node_memory, Memory.from_bytes(1)) # assert assert len(filtered_cycles) == 1 assert len(filtered_cycles[0]) == 2 assert set(n for n in filtered_cycles[0]) == {node1_id, node2_id} def test_filter_cycles_by_insufficient_memory(): # arrange node1_id = NodeId() node2_id = NodeId() connection1 = Connection( source=node1_id, sink=node2_id, edge=create_socket_connection(1) ) connection2 = Connection( source=node2_id, sink=node1_id, edge=create_socket_connection(2) ) node1_mem = create_node_memory(1000 * 1024) node2_mem = create_node_memory(1000 * 1024) node_memory = {node1_id: node1_mem, node2_id: node2_mem} topology = Topology() topology.add_node(node1_id) topology.add_node(node2_id) topology.add_connection(connection1) topology.add_connection(connection2) # act filtered_cycles = filter_cycles_by_memory( topology.get_cycles(), node_memory, Memory.from_kb(2001) ) # assert assert len(filtered_cycles) == 0 def test_filter_multiple_cycles_by_memory(): # arrange node_a_id = NodeId() node_b_id = NodeId() node_c_id = NodeId() connection1 = Connection( source=node_a_id, sink=node_b_id, edge=create_socket_connection(1) ) connection2 = Connection( source=node_b_id, sink=node_a_id, edge=create_socket_connection(2) ) connection3 = Connection( source=node_a_id, sink=node_c_id, edge=create_socket_connection(3) ) connection4 = Connection( source=node_c_id, sink=node_b_id, edge=create_socket_connection(4) ) node_a_mem = create_node_memory(500 * 1024) node_b_mem = create_node_memory(500 * 1024) node_c_mem = create_node_memory(1000 * 1024) node_memory = { node_a_id: node_a_mem, node_b_id: node_b_mem, node_c_id: node_c_mem, } topology = Topology() topology.add_node(node_a_id) topology.add_node(node_b_id) topology.add_node(node_c_id) topology.add_connection(connection1) topology.add_connection(connection2) topology.add_connection(connection3) topology.add_connection(connection4) cycles = topology.get_cycles() # act filtered_cycles = filter_cycles_by_memory(cycles, node_memory, Memory.from_kb(1500)) # assert assert len(filtered_cycles) == 1 assert len(filtered_cycles[0]) == 3 assert set(n for n in filtered_cycles[0]) == { node_a_id, node_b_id, node_c_id, } def test_get_smallest_cycles(): # arrange node_a_id = NodeId() node_b_id = NodeId() node_c_id = NodeId() topology = Topology() topology.add_node(node_a_id) topology.add_node(node_b_id) topology.add_node(node_c_id) connection1 = Connection( source=node_a_id, sink=node_b_id, edge=create_socket_connection(1) ) connection2 = Connection( source=node_b_id, sink=node_a_id, edge=create_socket_connection(2) ) connection3 = Connection( source=node_a_id, sink=node_c_id, edge=create_socket_connection(3) ) connection4 = Connection( source=node_c_id, sink=node_b_id, edge=create_socket_connection(4) ) topology.add_connection(connection1) topology.add_connection(connection2) topology.add_connection(connection3) topology.add_connection(connection4) cycles = [c for c in topology.get_cycles() if len(c) != 1] # ignore singletons # act smallest_cycles = get_smallest_cycles(cycles) # assert assert len(smallest_cycles) == 1 assert len(smallest_cycles[0]) == 2 assert set(n for n in smallest_cycles[0]) == {node_a_id, node_b_id} @pytest.mark.parametrize( "available_memory,total_layers,expected_layers", [ ((500, 500, 1000), 12, (3, 3, 6)), ((500, 500, 500), 12, (4, 4, 4)), ((312, 518, 1024), 12, (2, 3, 7)), # Edge case: one node has ~90% of memory - should not over-allocate. # Each node must have enough memory for at least 1 layer (50 KB = 1000/20). ((900, 50, 50), 20, (18, 1, 1)), ], ) def test_get_shard_assignments( available_memory: tuple[int, int, int], total_layers: int, expected_layers: tuple[int, int, int], ): # arrange node_a_id = NodeId() node_b_id = NodeId() node_c_id = NodeId() # create connections (A -> B -> C -> A forms a 3-cycle, plus B -> A also exists) connection1 = Connection( source=node_a_id, sink=node_b_id, edge=create_socket_connection(1) ) connection2 = Connection( source=node_b_id, sink=node_c_id, edge=create_socket_connection(2) ) connection3 = Connection( source=node_c_id, sink=node_a_id, edge=create_socket_connection(3) ) connection4 = Connection( source=node_b_id, sink=node_a_id, edge=create_socket_connection(4) ) topology = Topology() topology.add_node(node_a_id) topology.add_node(node_b_id) topology.add_node(node_c_id) topology.add_connection(connection1) topology.add_connection(connection2) topology.add_connection(connection3) topology.add_connection(connection4) node_a_mem = create_node_memory(available_memory[0] * 1024) node_b_mem = create_node_memory(available_memory[1] * 1024) node_c_mem = create_node_memory(available_memory[2] * 1024) node_memory = { node_a_id: node_a_mem, node_b_id: node_b_mem, node_c_id: node_c_mem, } model_card = ModelCard( model_id=ModelId("test-model"), n_layers=total_layers, storage_size=Memory.from_kb(1000), hidden_size=1000, supports_tensor=True, tasks=[ModelTask.TextGeneration], ) cycles = topology.get_cycles() # pick the 3-node cycle deterministically (cycle ordering can vary) selected_cycle = next(cycle for cycle in cycles if len(cycle) == 3) # act shard_assignments = get_shard_assignments( model_card, selected_cycle, Sharding.Pipeline, node_memory=node_memory ) # assert runner_id_a = shard_assignments.node_to_runner[node_a_id] runner_id_b = shard_assignments.node_to_runner[node_b_id] runner_id_c = shard_assignments.node_to_runner[node_c_id] assert ( shard_assignments.runner_to_shard[runner_id_a].end_layer - shard_assignments.runner_to_shard[runner_id_a].start_layer == expected_layers[0] ) assert ( shard_assignments.runner_to_shard[runner_id_b].end_layer - shard_assignments.runner_to_shard[runner_id_b].start_layer == expected_layers[1] ) assert ( shard_assignments.runner_to_shard[runner_id_c].end_layer - shard_assignments.runner_to_shard[runner_id_c].start_layer == expected_layers[2] ) def test_get_mlx_jaccl_coordinators(): # arrange node_a_id = NodeId() node_b_id = NodeId() node_c_id = NodeId() # fully connected (directed) between the 3 nodes conn_a_b = Connection( source=node_a_id, sink=node_b_id, edge=create_socket_connection(1) ) conn_b_a = Connection( source=node_b_id, sink=node_a_id, edge=create_socket_connection(2) ) conn_b_c = Connection( source=node_b_id, sink=node_c_id, edge=create_socket_connection(3) ) conn_c_b = Connection( source=node_c_id, sink=node_b_id, edge=create_socket_connection(4) ) conn_c_a = Connection( source=node_c_id, sink=node_a_id, edge=create_socket_connection(5) ) conn_a_c = Connection( source=node_a_id, sink=node_c_id, edge=create_socket_connection(6) ) network_a = NodeNetworkInfo( interfaces=[ NetworkInterfaceInfo(name="en0", ip_address="169.254.0.5"), NetworkInterfaceInfo(name="en0", ip_address="169.254.0.2"), ] ) network_b = NodeNetworkInfo( interfaces=[ NetworkInterfaceInfo(name="en0", ip_address="169.254.0.1"), NetworkInterfaceInfo(name="en0", ip_address="169.254.0.4"), ] ) network_c = NodeNetworkInfo( interfaces=[ NetworkInterfaceInfo(name="en0", ip_address="169.254.0.3"), NetworkInterfaceInfo(name="en0", ip_address="169.254.0.6"), ] ) node_network = { node_a_id: network_a, node_b_id: network_b, node_c_id: network_c, } topology = Topology() topology.add_node(node_a_id) topology.add_node(node_b_id) topology.add_node(node_c_id) topology.add_connection(conn_a_b) topology.add_connection(conn_b_a) topology.add_connection(conn_b_c) topology.add_connection(conn_c_b) topology.add_connection(conn_c_a) topology.add_connection(conn_a_c) # act coordinators = get_mlx_jaccl_coordinators( node_a_id, coordinator_port=5000, cycle_digraph=topology, node_network=node_network, ) # assert assert len(coordinators) == 3 assert node_a_id in coordinators assert node_b_id in coordinators assert node_c_id in coordinators # All coordinators should have IP:PORT format for node_id, coordinator in coordinators.items(): assert ":" in coordinator, ( f"Coordinator for {node_id} should have ':' separator" ) # Verify port is correct for node_id, coordinator in coordinators.items(): assert coordinator.endswith(":5000"), ( f"Coordinator for {node_id} should use port 5000" ) # Rank 0 (node_a) treats this as the listen socket so should listen on all IPs assert coordinators[node_a_id].startswith("0.0.0.0:"), ( "Rank 0 node should use 0.0.0.0 as coordinator listen address" ) # Non-rank-0 nodes should use the specific IP from their connection to rank 0 # node_b uses the IP from conn_b_a (node_b -> node_a) assert isinstance(conn_b_a.edge, SocketConnection) assert ( coordinators[node_b_id] == f"{conn_b_a.edge.sink_multiaddr.ip_address}:5000" ), "node_b should use the IP from conn_b_a" # node_c uses the IP from conn_c_a (node_c -> node_a) assert isinstance(conn_c_a.edge, SocketConnection) assert coordinators[node_c_id] == ( f"{conn_c_a.edge.sink_multiaddr.ip_address}:5000" ), "node_c should use the IP from conn_c_a" class TestAllocateLayersProportionally: def test_empty_node_list_raises(self): with pytest.raises(ValueError, match="empty node list"): allocate_layers_proportionally(total_layers=10, memory_fractions=[]) def test_zero_layers_raises(self): with pytest.raises(ValueError, match="need at least 1 layer per node"): allocate_layers_proportionally(total_layers=0, memory_fractions=[0.5, 0.5]) def test_negative_layers_raises(self): with pytest.raises(ValueError, match="need at least 1 layer per node"): allocate_layers_proportionally(total_layers=-1, memory_fractions=[0.5, 0.5]) def test_fewer_layers_than_nodes_raises(self): with pytest.raises(ValueError, match="need at least 1 layer per node"): allocate_layers_proportionally( total_layers=2, memory_fractions=[0.33, 0.33, 0.34] ) def test_equal_distribution(self): result = allocate_layers_proportionally( total_layers=12, memory_fractions=[0.25, 0.25, 0.25, 0.25] ) assert result == [3, 3, 3, 3] assert sum(result) == 12 def test_proportional_distribution(self): result = allocate_layers_proportionally( total_layers=12, memory_fractions=[0.25, 0.25, 0.50] ) assert result == [3, 3, 6] assert sum(result) == 12 def test_extreme_imbalance_ensures_minimum(self): result = allocate_layers_proportionally( total_layers=20, memory_fractions=[0.975, 0.0125, 0.0125] ) assert all(layers >= 1 for layers in result) assert sum(result) == 20 # Small nodes get minimum 1 layer assert result == [18, 1, 1] def test_single_node_gets_all_layers(self): result = allocate_layers_proportionally(total_layers=10, memory_fractions=[1.0]) assert result == [10] def test_minimum_viable_allocation(self): result = allocate_layers_proportionally( total_layers=3, memory_fractions=[0.33, 0.33, 0.34] ) assert result == [1, 1, 1] assert sum(result) == 3 def test_get_shard_assignments_insufficient_memory_raises(): """Test that ValueError is raised when a node has insufficient memory for its layers.""" node_a_id = NodeId() node_b_id = NodeId() node_c_id = NodeId() topology = Topology() # Node C has only 10 KB but would need 50 KB for 1 layer (1000 KB / 20 layers) node_a_mem = create_node_memory(900 * 1024) node_b_mem = create_node_memory(50 * 1024) node_c_mem = create_node_memory(10 * 1024) # Insufficient memory topology.add_node(node_a_id) topology.add_node(node_b_id) topology.add_node(node_c_id) conn_a_b = Connection( source=node_a_id, sink=node_b_id, edge=create_socket_connection(1) ) conn_b_c = Connection( source=node_b_id, sink=node_c_id, edge=create_socket_connection(2) ) conn_c_a = Connection( source=node_c_id, sink=node_a_id, edge=create_socket_connection(3) ) conn_b_a = Connection( source=node_b_id, sink=node_a_id, edge=create_socket_connection(3) ) topology.add_connection(conn_a_b) topology.add_connection(conn_b_c) topology.add_connection(conn_c_a) topology.add_connection(conn_b_a) node_memory = { node_a_id: node_a_mem, node_b_id: node_b_mem, node_c_id: node_c_mem, } model_card = ModelCard( model_id=ModelId("test-model"), n_layers=20, storage_size=Memory.from_kb(1000), hidden_size=1000, supports_tensor=True, tasks=[ModelTask.TextGeneration], ) cycles = topology.get_cycles() selected_cycle = cycles[0] with pytest.raises(ValueError, match="insufficient memory"): get_shard_assignments( model_card, selected_cycle, Sharding.Pipeline, node_memory ) class TestCfgParallelPlacement: def _create_ring_topology(self, node_ids: list[NodeId]) -> Topology: topology = Topology() for node_id in node_ids: topology.add_node(node_id) for i, node_id in enumerate(node_ids): next_node = node_ids[(i + 1) % len(node_ids)] conn = Connection( source=node_id, sink=next_node, edge=create_socket_connection(i + 1), ) topology.add_connection(conn) return topology def test_two_nodes_cfg_model_uses_cfg_parallel(self): """Two nodes with CFG model should use CFG parallel (no pipeline).""" node_a = NodeId() node_b = NodeId() topology = self._create_ring_topology([node_a, node_b]) cycles = [c for c in topology.get_cycles() if len(c) == 2] cycle = cycles[0] node_memory = { node_a: create_node_memory(1000 * 1024), node_b: create_node_memory(1000 * 1024), } model_card = ModelCard( model_id=ModelId("qwen-image-test"), n_layers=60, storage_size=Memory.from_kb(1000), hidden_size=1, supports_tensor=False, uses_cfg=True, tasks=[ModelTask.TextToImage], ) assignments = get_shard_assignments_for_pipeline_parallel( model_card, cycle, node_memory ) shards = list(assignments.runner_to_shard.values()) assert len(shards) == 2 # CFG models should get CfgShardMetadata for shard in shards: assert isinstance(shard, CfgShardMetadata) # Both nodes should have all layers (no pipeline split) assert shard.start_layer == 0 assert shard.end_layer == 60 assert shard.cfg_world_size == 2 # Each node is the only stage in its pipeline group assert shard.pipeline_world_size == 1 assert shard.pipeline_rank == 0 cfg_ranks = sorted( s.cfg_rank for s in shards if isinstance(s, CfgShardMetadata) ) assert cfg_ranks == [0, 1] def test_four_nodes_cfg_model_uses_hybrid(self): """Four nodes with CFG model should use 2 CFG groups x 2 pipeline stages.""" nodes = [NodeId() for _ in range(4)] topology = self._create_ring_topology(nodes) cycles = [c for c in topology.get_cycles() if len(c) == 4] cycle = cycles[0] node_memory = {n: create_node_memory(1000 * 1024) for n in nodes} model_card = ModelCard( model_id=ModelId("qwen-image-test"), n_layers=60, storage_size=Memory.from_kb(1000), hidden_size=1, supports_tensor=False, uses_cfg=True, tasks=[ModelTask.TextToImage], ) assignments = get_shard_assignments_for_pipeline_parallel( model_card, cycle, node_memory ) shards = list(assignments.runner_to_shard.values()) assert len(shards) == 4 # CFG models should get CfgShardMetadata for shard in shards: assert isinstance(shard, CfgShardMetadata) assert shard.cfg_world_size == 2 assert shard.pipeline_world_size == 2 assert shard.pipeline_rank in [0, 1] # Check we have 2 nodes in each CFG group cfg_0_shards = [ s for s in shards if isinstance(s, CfgShardMetadata) and s.cfg_rank == 0 ] cfg_1_shards = [ s for s in shards if isinstance(s, CfgShardMetadata) and s.cfg_rank == 1 ] assert len(cfg_0_shards) == 2 assert len(cfg_1_shards) == 2 # Both CFG groups should have the same layer assignments cfg_0_layers = [(s.start_layer, s.end_layer) for s in cfg_0_shards] cfg_1_layers = [(s.start_layer, s.end_layer) for s in cfg_1_shards] assert sorted(cfg_0_layers) == sorted(cfg_1_layers) def test_three_nodes_cfg_model_uses_sequential_cfg(self): """Three nodes (odd) with CFG model should use sequential CFG (PipelineShardMetadata).""" nodes = [NodeId() for _ in range(3)] topology = self._create_ring_topology(nodes) cycles = [c for c in topology.get_cycles() if len(c) == 3] cycle = cycles[0] node_memory = {n: create_node_memory(1000 * 1024) for n in nodes} model_card = ModelCard( model_id=ModelId("qwen-image-test"), n_layers=60, storage_size=Memory.from_kb(1000), hidden_size=1, supports_tensor=False, uses_cfg=True, tasks=[ModelTask.TextToImage], ) assignments = get_shard_assignments_for_pipeline_parallel( model_card, cycle, node_memory ) shards = list(assignments.runner_to_shard.values()) assert len(shards) == 3 # Odd node count with CFG model falls back to PipelineShardMetadata (sequential CFG) for shard in shards: assert isinstance(shard, PipelineShardMetadata) def test_two_nodes_non_cfg_model_uses_pipeline(self): """Two nodes with non-CFG model should use pure pipeline (PipelineShardMetadata).""" node_a = NodeId() node_b = NodeId() topology = self._create_ring_topology([node_a, node_b]) cycles = [c for c in topology.get_cycles() if len(c) == 2] cycle = cycles[0] node_memory = { node_a: create_node_memory(1000 * 1024), node_b: create_node_memory(1000 * 1024), } model_card = ModelCard( model_id=ModelId("flux-test"), n_layers=57, storage_size=Memory.from_kb(1000), hidden_size=1, supports_tensor=False, uses_cfg=False, # Non-CFG model tasks=[ModelTask.TextToImage], ) assignments = get_shard_assignments_for_pipeline_parallel( model_card, cycle, node_memory ) shards = list(assignments.runner_to_shard.values()) assert len(shards) == 2 # Non-CFG models should get PipelineShardMetadata for shard in shards: assert isinstance(shard, PipelineShardMetadata) # Should have actual layer sharding (pipeline) layer_ranges = sorted( (s.start_layer, s.end_layer) for s in shards if isinstance(s, PipelineShardMetadata) ) # First shard starts at 0, last shard ends at 57 assert layer_ranges[0][0] == 0 assert layer_ranges[-1][1] == 57 ================================================ FILE: src/exo/master/tests/test_topology.py ================================================ import pytest from exo.shared.topology import Topology from exo.shared.types.common import NodeId from exo.shared.types.multiaddr import Multiaddr from exo.shared.types.topology import Connection, SocketConnection @pytest.fixture def topology() -> Topology: return Topology() @pytest.fixture def socket_connection() -> SocketConnection: return SocketConnection( sink_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/1235"), ) def test_add_node(topology: Topology): # arrange node_id = NodeId() # act topology.add_node(node_id) # assert assert topology.node_is_leaf(node_id) def test_add_connection(topology: Topology, socket_connection: SocketConnection): # arrange node_a = NodeId() node_b = NodeId() connection = Connection(source=node_a, sink=node_b, edge=socket_connection) topology.add_node(node_a) topology.add_node(node_b) topology.add_connection(connection) # act data = list(topology.list_connections()) # assert assert data == [connection] assert topology.node_is_leaf(node_a) assert topology.node_is_leaf(node_b) def test_remove_connection_still_connected( topology: Topology, socket_connection: SocketConnection ): # arrange node_a = NodeId() node_b = NodeId() conn = Connection(source=node_a, sink=node_b, edge=socket_connection) topology.add_node(node_a) topology.add_node(node_b) topology.add_connection(conn) # act topology.remove_connection(conn) # assert assert list(topology.get_all_connections_between(node_a, node_b)) == [] def test_remove_node_still_connected( topology: Topology, socket_connection: SocketConnection ): # arrange node_a = NodeId() node_b = NodeId() conn = Connection(source=node_a, sink=node_b, edge=socket_connection) topology.add_node(node_a) topology.add_node(node_b) topology.add_connection(conn) assert list(topology.out_edges(node_a)) == [conn] # act topology.remove_node(node_b) # assert assert list(topology.out_edges(node_a)) == [] def test_list_nodes(topology: Topology, socket_connection: SocketConnection): # arrange node_a = NodeId() node_b = NodeId() conn = Connection(source=node_a, sink=node_b, edge=socket_connection) topology.add_node(node_a) topology.add_node(node_b) topology.add_connection(conn) assert list(topology.out_edges(node_a)) == [conn] # act nodes = list(topology.list_nodes()) # assert assert len(nodes) == 2 assert all(isinstance(node, NodeId) for node in nodes) assert set(node for node in nodes) == set([node_a, node_b]) ================================================ FILE: src/exo/routing/__init__.py ================================================ ================================================ FILE: src/exo/routing/connection_message.py ================================================ from exo_pyo3_bindings import PyFromSwarm from exo.shared.types.common import NodeId from exo.utils.pydantic_ext import CamelCaseModel """Serialisable types for Connection Updates/Messages""" class ConnectionMessage(CamelCaseModel): node_id: NodeId connected: bool @classmethod def from_update(cls, update: PyFromSwarm.Connection) -> "ConnectionMessage": return cls(node_id=NodeId(update.peer_id), connected=update.connected) ================================================ FILE: src/exo/routing/event_router.py ================================================ from dataclasses import dataclass, field from random import random import anyio from anyio import BrokenResourceError, ClosedResourceError from anyio.abc import CancelScope from loguru import logger from exo.shared.types.commands import ForwarderCommand, RequestEventLog from exo.shared.types.common import SessionId, SystemId from exo.shared.types.events import ( Event, EventId, GlobalForwarderEvent, IndexedEvent, LocalForwarderEvent, ) from exo.utils.channels import Receiver, Sender, channel from exo.utils.event_buffer import OrderedBuffer from exo.utils.task_group import TaskGroup @dataclass class EventRouter: session_id: SessionId command_sender: Sender[ForwarderCommand] external_inbound: Receiver[GlobalForwarderEvent] external_outbound: Sender[LocalForwarderEvent] _system_id: SystemId = field(init=False, default_factory=SystemId) internal_outbound: list[Sender[IndexedEvent]] = field( init=False, default_factory=list ) event_buffer: OrderedBuffer[Event] = field( init=False, default_factory=OrderedBuffer ) out_for_delivery: dict[EventId, tuple[float, LocalForwarderEvent]] = field( init=False, default_factory=dict ) _tg: TaskGroup = field(init=False, default_factory=TaskGroup) _nack_cancel_scope: CancelScope | None = field(init=False, default=None) _nack_attempts: int = field(init=False, default=0) _nack_base_seconds: float = field(init=False, default=0.5) _nack_cap_seconds: float = field(init=False, default=10.0) async def run(self): try: async with self._tg as tg: tg.start_soon(self._run_ext_in) tg.start_soon(self._simple_retry) finally: self.external_outbound.close() for send in self.internal_outbound: send.close() # can make this better in future async def _simple_retry(self): while True: await anyio.sleep(1 + random()) # list here is a shallow clone for shared mutation for e_id, (time, event) in list(self.out_for_delivery.items()): if anyio.current_time() > time + 5: self.out_for_delivery[e_id] = (anyio.current_time(), event) await self.external_outbound.send(event) def sender(self) -> Sender[Event]: send, recv = channel[Event]() if self._tg.is_running(): self._tg.start_soon(self._ingest, SystemId(), recv) else: self._tg.queue(self._ingest, SystemId(), recv) return send def receiver(self) -> Receiver[IndexedEvent]: send, recv = channel[IndexedEvent]() self.internal_outbound.append(send) return recv def shutdown(self) -> None: self._tg.cancel_tasks() async def _ingest(self, system_id: SystemId, recv: Receiver[Event]): idx = 0 with recv as events: async for event in events: f_ev = LocalForwarderEvent( origin_idx=idx, origin=system_id, session=self.session_id, event=event, ) idx += 1 await self.external_outbound.send(f_ev) self.out_for_delivery[event.event_id] = (anyio.current_time(), f_ev) async def _run_ext_in(self): buf = OrderedBuffer[Event]() with self.external_inbound as events: async for event in events: if event.session != self.session_id: continue if event.origin != self.session_id.master_node_id: continue buf.ingest(event.origin_idx, event.event) event_id = event.event.event_id if event_id in self.out_for_delivery: self.out_for_delivery.pop(event_id) drained = buf.drain_indexed() if drained: self._nack_attempts = 0 if self._nack_cancel_scope: self._nack_cancel_scope.cancel() if not drained and ( self._nack_cancel_scope is None or self._nack_cancel_scope.cancel_called ): # Request the next index. self._tg.start_soon(self._nack_request, buf.next_idx_to_release) continue for idx, event in drained: to_clear = set[int]() for i, sender in enumerate(self.internal_outbound): try: await sender.send(IndexedEvent(idx=idx, event=event)) except (ClosedResourceError, BrokenResourceError): to_clear.add(i) for i in sorted(to_clear, reverse=True): self.internal_outbound.pop(i) async def _nack_request(self, since_idx: int) -> None: # We request all events after (and including) the missing index. # This function is started whenever we receive an event that is out of sequence. # It is cancelled as soon as we receiver an event that is in sequence. if since_idx < 0: logger.warning(f"Negative value encountered for nack request {since_idx=}") since_idx = 0 with CancelScope() as scope: self._nack_cancel_scope = scope delay: float = self._nack_base_seconds * (2.0**self._nack_attempts) delay = min(self._nack_cap_seconds, delay) self._nack_attempts += 1 try: await anyio.sleep(delay) logger.info( f"Nack attempt {self._nack_attempts}: Requesting Event Log from {since_idx}" ) await self.command_sender.send( ForwarderCommand( origin=self._system_id, command=RequestEventLog(since_idx=since_idx), ) ) finally: if self._nack_cancel_scope is scope: self._nack_cancel_scope = None ================================================ FILE: src/exo/routing/router.py ================================================ from copy import copy from itertools import count from math import inf from os import PathLike from pathlib import Path from typing import cast from anyio import ( BrokenResourceError, ClosedResourceError, move_on_after, sleep_forever, ) from exo_pyo3_bindings import ( AllQueuesFullError, Keypair, MessageTooLargeError, NetworkingHandle, NoPeersSubscribedToTopicError, PyFromSwarm, ) from filelock import FileLock from loguru import logger from exo.shared.constants import EXO_NODE_ID_KEYPAIR from exo.utils.channels import Receiver, Sender, channel from exo.utils.pydantic_ext import CamelCaseModel from exo.utils.task_group import TaskGroup from .connection_message import ConnectionMessage from .topics import CONNECTION_MESSAGES, PublishPolicy, TypedTopic # A significant current limitation of the TopicRouter is that it is not capable # of preventing feedback, as it does not ask for a system id so cannot tell # which message is coming/going to which system. # This is currently only relevant for Election class TopicRouter[T: CamelCaseModel]: def __init__( self, topic: TypedTopic[T], networking_sender: Sender[tuple[str, bytes]], max_buffer_size: float = inf, ): self.topic: TypedTopic[T] = topic self.senders: set[Sender[T]] = set() send, recv = channel[T]() self.receiver: Receiver[T] = recv self._sender: Sender[T] = send self.networking_sender: Sender[tuple[str, bytes]] = networking_sender async def run(self): logger.debug(f"Topic Router {self.topic} ready to send") with self.receiver as items: async for item in items: # Check if we should send to network if ( len(self.senders) == 0 and self.topic.publish_policy is PublishPolicy.Minimal ): await self._send_out(item) continue if self.topic.publish_policy is PublishPolicy.Always: await self._send_out(item) # Then publish to all senders await self.publish(item) async def shutdown(self): logger.debug(f"Shutting down Topic Router {self.topic}") # Close all the things! for sender in self.senders: sender.close() self._sender.close() self.receiver.close() async def publish(self, item: T): """ Publish item T on this topic to all senders. NB: this sends to ALL receivers, potentially including receivers held by the object doing the sending. You should handle your own output if you hold a sender + receiver pair. """ to_clear: set[Sender[T]] = set() for sender in copy(self.senders): try: await sender.send(item) except (ClosedResourceError, BrokenResourceError): to_clear.add(sender) self.senders -= to_clear async def publish_bytes(self, data: bytes): await self.publish(self.topic.deserialize(data)) def new_sender(self) -> Sender[T]: return self._sender.clone() async def _send_out(self, item: T): logger.trace(f"TopicRouter {self.topic.topic} sending {item}") await self.networking_sender.send( (str(self.topic.topic), self.topic.serialize(item)) ) class Router: @classmethod def create(cls, identity: Keypair) -> "Router": return cls(handle=NetworkingHandle(identity)) def __init__(self, handle: NetworkingHandle): self.topic_routers: dict[str, TopicRouter[CamelCaseModel]] = {} send, recv = channel[tuple[str, bytes]]() self.networking_receiver: Receiver[tuple[str, bytes]] = recv self._net: NetworkingHandle = handle self._tmp_networking_sender: Sender[tuple[str, bytes]] | None = send self._id_count = count() self._tg: TaskGroup = TaskGroup() async def register_topic[T: CamelCaseModel](self, topic: TypedTopic[T]): send = self._tmp_networking_sender if send: self._tmp_networking_sender = None else: send = self.networking_receiver.clone_sender() router = TopicRouter[T](topic, send) self.topic_routers[topic.topic] = cast(TopicRouter[CamelCaseModel], router) if self._tg.is_running(): await self._networking_subscribe(topic.topic) def sender[T: CamelCaseModel](self, topic: TypedTopic[T]) -> Sender[T]: router = self.topic_routers.get(topic.topic, None) # There's gotta be a way to do this without THIS many asserts assert router is not None assert router.topic == topic sender = cast(TopicRouter[T], router).new_sender() return sender def receiver[T: CamelCaseModel](self, topic: TypedTopic[T]) -> Receiver[T]: router = self.topic_routers.get(topic.topic, None) # There's gotta be a way to do this without THIS many asserts assert router is not None assert router.topic == topic assert router.topic.model_type == topic.model_type send, recv = channel[T]() router.senders.add(cast(Sender[CamelCaseModel], send)) return recv async def run(self): logger.debug("Starting Router") try: async with self._tg as tg: for topic in self.topic_routers: router = self.topic_routers[topic] tg.start_soon(router.run) tg.start_soon(self._networking_recv) tg.start_soon(self._networking_publish) # subscribe to pending topics for topic in self.topic_routers: await self._networking_subscribe(topic) # Router only shuts down if you cancel it. await sleep_forever() finally: with move_on_after(1, shield=True): for topic in self.topic_routers: await self._networking_unsubscribe(str(topic)) async def shutdown(self): logger.debug("Shutting down Router") self._tg.cancel_tasks() async def _networking_subscribe(self, topic: str): await self._net.gossipsub_subscribe(topic) logger.info(f"Subscribed to {topic}") async def _networking_unsubscribe(self, topic: str): await self._net.gossipsub_unsubscribe(topic) logger.info(f"Unsubscribed from {topic}") async def _networking_recv(self): try: while True: from_swarm = await self._net.recv() logger.debug(from_swarm) match from_swarm: case PyFromSwarm.Message(origin, topic, data): logger.trace( f"Received message on {topic} from {origin} with payload {data}" ) if topic not in self.topic_routers: logger.warning( f"Received message on unknown or inactive topic {topic}" ) continue router = self.topic_routers[topic] await router.publish_bytes(data) case PyFromSwarm.Connection(): message = ConnectionMessage.from_update(from_swarm) logger.trace( f"Received message on connection_messages with payload {message}" ) if CONNECTION_MESSAGES.topic in self.topic_routers: router = self.topic_routers[CONNECTION_MESSAGES.topic] assert router.topic.model_type == ConnectionMessage router = cast(TopicRouter[ConnectionMessage], router) await router.publish(message) case _: logger.critical( "failed to exhaustively check FromSwarm messages - logic error" ) except Exception as exception: logger.opt(exception=exception).error( "Gossipsub receive loop terminated unexpectedly" ) raise async def _networking_publish(self): with self.networking_receiver as networked_items: async for topic, data in networked_items: try: logger.trace(f"Sending message on {topic} with payload {data}") if len(data) > 1024 * 1024: logger.warning( "Sending overlarge payload, network performance may be temporarily degraded" ) await self._net.gossipsub_publish(topic, data) except NoPeersSubscribedToTopicError: pass except AllQueuesFullError: logger.warning(f"All peer queues full, dropping message on {topic}") except MessageTooLargeError: logger.warning( f"Message too large for gossipsub on {topic} ({len(data)} bytes), dropping" ) def get_node_id_keypair( path: str | bytes | PathLike[str] | PathLike[bytes] = EXO_NODE_ID_KEYPAIR, ) -> Keypair: """ Obtains the :class:`Keypair` associated with this node-ID. Obtain the :class:`PeerId` by from it. """ # TODO(evan): bring back node id persistence once we figure out how to deal with duplicates return Keypair.generate() def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path: return Path(str(path) + ".lock") # operate with cross-process lock to avoid race conditions with FileLock(lock_path(path)): with open(path, "a+b") as f: # opens in append-mode => starts at EOF # if non-zero EOF, then file exists => use to get node-ID if f.tell() != 0: f.seek(0) # go to start & read protobuf-encoded bytes protobuf_encoded = f.read() try: # if decoded successfully, save & return return Keypair.from_bytes(protobuf_encoded) except ValueError as e: # on runtime error, assume corrupt file logger.warning(f"Encountered error when trying to get keypair: {e}") # if no valid credentials, create new ones and persist with open(path, "w+b") as f: keypair = Keypair.generate() f.write(keypair.to_bytes()) return keypair ================================================ FILE: src/exo/routing/tests/test_event_buffer.py ================================================ import pytest from exo.shared.types.events import Event, TestEvent from exo.utils.event_buffer import OrderedBuffer def make_indexed_event(idx: int) -> tuple[int, Event]: """Factory function to create a unique ForwarderEvent for a given index.""" return (idx, TestEvent()) @pytest.fixture def buffer() -> OrderedBuffer[Event]: """Provides a clean instance of OrderedBuffer[Event] for each test.""" return OrderedBuffer[Event]() @pytest.mark.asyncio async def test_initial_state(buffer: OrderedBuffer[Event]): """Tests that a new buffer is empty and starts at index 1.""" assert buffer.next_idx_to_release == 0 assert not buffer.store assert buffer.drain() == [] @pytest.mark.asyncio async def test_ingest_and_drain_sequential_events(buffer: OrderedBuffer[Event]): """Tests ingesting and draining a simple, ordered sequence of events.""" events = [make_indexed_event(0), make_indexed_event(1), make_indexed_event(2)] [buffer.ingest(*ev) for ev in events] drained_events = buffer.drain_indexed() assert drained_events == events assert buffer.next_idx_to_release == 3 assert not buffer.store @pytest.mark.asyncio async def test_ingest_out_of_order_events(buffer: OrderedBuffer[Event]): """Tests that out-of-order events are buffered and drained in the correct sequence.""" event1 = make_indexed_event(0) event2 = make_indexed_event(1) event3 = make_indexed_event(2) buffer.ingest(*event3) buffer.ingest(*event1) buffer.ingest(*event2) drained_events = buffer.drain_indexed() assert drained_events == [event1, event2, event3] assert buffer.next_idx_to_release == 3 @pytest.mark.asyncio async def test_drain_with_gap_in_sequence(buffer: OrderedBuffer[Event]): """Tests that draining stops when there is a gap in the event indices.""" event1 = make_indexed_event(0) event3 = make_indexed_event(2) buffer.ingest(*event1) buffer.ingest(*event3) drained_events = buffer.drain_indexed() assert drained_events == [event1] assert buffer.next_idx_to_release == 1 assert buffer.drain() == [] assert 2 in buffer.store @pytest.mark.asyncio async def test_fill_gap_and_drain_remaining(buffer: OrderedBuffer[Event]): """Tests that once a gap is filled, the rest of the sequence is drained.""" event0 = make_indexed_event(0) event2 = make_indexed_event(2) buffer.ingest(*event0) buffer.ingest(*event2) buffer.drain() assert buffer.next_idx_to_release == 1 event1 = make_indexed_event(1) buffer.ingest(*event1) drained_events = buffer.drain_indexed() assert [e[0] for e in drained_events] == [1, 2] assert buffer.next_idx_to_release == 3 @pytest.mark.asyncio async def test_ingest_drops_duplicate_indices(buffer: OrderedBuffer[Event]): """Tests that if multiple events for the same index are ingested, the first one wins.""" event2_first = make_indexed_event(1) event2_second = (1, TestEvent()) buffer.ingest(*make_indexed_event(0)) buffer.ingest(*event2_first) with pytest.raises(AssertionError): buffer.ingest(*event2_second) # This duplicate should be ignored drained = buffer.drain_indexed() assert len(drained) == 2 assert drained[1][1].event_id == event2_first[1].event_id assert drained[1][1].event_id != event2_second[1].event_id @pytest.mark.asyncio async def test_ingest_drops_stale_events(buffer: OrderedBuffer[Event]): """Tests that events with an index lower than next_idx_to_release are dropped.""" buffer.ingest(*make_indexed_event(0)) buffer.ingest(*make_indexed_event(1)) buffer.drain() assert buffer.next_idx_to_release == 2 stale_event1 = make_indexed_event(0) stale_event2 = make_indexed_event(1) buffer.ingest(*stale_event1) buffer.ingest(*stale_event2) assert not buffer.store assert buffer.drain() == [] @pytest.mark.asyncio async def test_drain_and_ingest_with_new_sequence(buffer: OrderedBuffer[Event]): """Tests reusing the buffer after it has been fully drained.""" buffer.ingest(*make_indexed_event(0)) buffer.ingest(*make_indexed_event(1)) buffer.drain() assert buffer.next_idx_to_release == 2 assert not buffer.store buffer.ingest(*make_indexed_event(4)) buffer.ingest(*make_indexed_event(2)) drained = buffer.drain_indexed() assert [e[0] for e in drained] == [2] assert buffer.next_idx_to_release == 3 assert 4 in buffer.store ================================================ FILE: src/exo/routing/topics.py ================================================ from dataclasses import dataclass from enum import Enum from exo.routing.connection_message import ConnectionMessage from exo.shared.election import ElectionMessage from exo.shared.types.commands import ForwarderCommand, ForwarderDownloadCommand from exo.shared.types.events import ( GlobalForwarderEvent, LocalForwarderEvent, ) from exo.utils.pydantic_ext import CamelCaseModel class PublishPolicy(str, Enum): Never = "Never" """Never publish to the network - this is a local message""" Minimal = "Minimal" """Only publish when there is no local receiver for this type of message""" Always = "Always" """Always publish to the network""" @dataclass # (frozen=True) class TypedTopic[T: CamelCaseModel]: topic: str publish_policy: PublishPolicy model_type: type[ T ] # This can be worked around with evil type hacking, see https://stackoverflow.com/a/71720366 - I don't think it's necessary here. @staticmethod def serialize(t: T) -> bytes: return t.model_dump_json().encode("utf-8") def deserialize(self, b: bytes) -> T: return self.model_type.model_validate_json(b.decode("utf-8")) GLOBAL_EVENTS = TypedTopic("global_events", PublishPolicy.Always, GlobalForwarderEvent) LOCAL_EVENTS = TypedTopic("local_events", PublishPolicy.Always, LocalForwarderEvent) COMMANDS = TypedTopic("commands", PublishPolicy.Always, ForwarderCommand) ELECTION_MESSAGES = TypedTopic( "election_messages", PublishPolicy.Always, ElectionMessage ) CONNECTION_MESSAGES = TypedTopic( "connection_messages", PublishPolicy.Never, ConnectionMessage ) DOWNLOAD_COMMANDS = TypedTopic( "download_commands", PublishPolicy.Always, ForwarderDownloadCommand ) ================================================ FILE: src/exo/shared/__init__.py ================================================ ================================================ FILE: src/exo/shared/apply.py ================================================ import copy from collections.abc import Mapping, Sequence from datetime import datetime from loguru import logger from exo.shared.types.common import NodeId from exo.shared.types.events import ( ChunkGenerated, Event, IndexedEvent, InputChunkReceived, InstanceCreated, InstanceDeleted, NodeDownloadProgress, NodeGatheredInfo, NodeTimedOut, RunnerStatusUpdated, TaskAcknowledged, TaskCreated, TaskDeleted, TaskFailed, TaskStatusUpdated, TestEvent, TopologyEdgeCreated, TopologyEdgeDeleted, TracesCollected, TracesMerged, ) from exo.shared.types.profiling import ( NodeIdentity, NodeNetworkInfo, NodeRdmaCtlStatus, NodeThunderboltInfo, ThunderboltBridgeStatus, ) from exo.shared.types.state import State from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.topology import Connection, RDMAConnection from exo.shared.types.worker.downloads import DownloadProgress from exo.shared.types.worker.instances import Instance, InstanceId from exo.shared.types.worker.runners import RunnerId, RunnerShutdown, RunnerStatus from exo.utils.info_gatherer.info_gatherer import ( MacmonMetrics, MacThunderboltConnections, MacThunderboltIdentifiers, MemoryUsage, MiscData, NodeConfig, NodeDiskUsage, NodeNetworkInterfaces, RdmaCtlStatus, StaticNodeInformation, ThunderboltBridgeInfo, ) def event_apply(event: Event, state: State) -> State: """Apply an event to state.""" match event: case ( TestEvent() | ChunkGenerated() | TaskAcknowledged() | InputChunkReceived() | TracesCollected() | TracesMerged() ): # Pass-through events that don't modify state return state case InstanceCreated(): return apply_instance_created(event, state) case InstanceDeleted(): return apply_instance_deleted(event, state) case NodeTimedOut(): return apply_node_timed_out(event, state) case NodeDownloadProgress(): return apply_node_download_progress(event, state) case NodeGatheredInfo(): return apply_node_gathered_info(event, state) case RunnerStatusUpdated(): return apply_runner_status_updated(event, state) case TaskCreated(): return apply_task_created(event, state) case TaskDeleted(): return apply_task_deleted(event, state) case TaskFailed(): return apply_task_failed(event, state) case TaskStatusUpdated(): return apply_task_status_updated(event, state) case TopologyEdgeCreated(): return apply_topology_edge_created(event, state) case TopologyEdgeDeleted(): return apply_topology_edge_deleted(event, state) def apply(state: State, event: IndexedEvent) -> State: # Just to test that events are only applied in correct order if state.last_event_applied_idx != event.idx - 1: logger.warning( f"Expected event {state.last_event_applied_idx + 1} but received {event.idx}" ) assert state.last_event_applied_idx == event.idx - 1 new_state: State = event_apply(event.event, state) return new_state.model_copy(update={"last_event_applied_idx": event.idx}) def apply_node_download_progress(event: NodeDownloadProgress, state: State) -> State: """ Update or add a node download progress to state. """ dp = event.download_progress node_id = dp.node_id current = list(state.downloads.get(node_id, ())) replaced = False for i, existing_dp in enumerate(current): # TODO(ciaran): deduplicate by model_id for now. Will need to use # shard_metadata again when pipeline and tensor downloads differ. # For now this is fine if ( existing_dp.shard_metadata.model_card.model_id == dp.shard_metadata.model_card.model_id ): current[i] = dp replaced = True break if not replaced: current.append(dp) new_downloads: Mapping[NodeId, Sequence[DownloadProgress]] = { **state.downloads, node_id: current, } return state.model_copy(update={"downloads": new_downloads}) def apply_task_created(event: TaskCreated, state: State) -> State: new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: event.task} return state.model_copy(update={"tasks": new_tasks}) def apply_task_deleted(event: TaskDeleted, state: State) -> State: new_tasks: Mapping[TaskId, Task] = { tid: task for tid, task in state.tasks.items() if tid != event.task_id } return state.model_copy(update={"tasks": new_tasks}) def apply_task_status_updated(event: TaskStatusUpdated, state: State) -> State: if event.task_id not in state.tasks: # maybe should raise return state update: dict[str, TaskStatus | None] = { "task_status": event.task_status, } if event.task_status != TaskStatus.Failed: update["error_type"] = None update["error_message"] = None updated_task = state.tasks[event.task_id].model_copy(update=update) new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: updated_task} return state.model_copy(update={"tasks": new_tasks}) def apply_task_failed(event: TaskFailed, state: State) -> State: if event.task_id not in state.tasks: # maybe should raise return state updated_task = state.tasks[event.task_id].model_copy( update={"error_type": event.error_type, "error_message": event.error_message} ) new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: updated_task} return state.model_copy(update={"tasks": new_tasks}) def apply_instance_created(event: InstanceCreated, state: State) -> State: instance = event.instance new_instances: Mapping[InstanceId, Instance] = { **state.instances, instance.instance_id: instance, } return state.model_copy(update={"instances": new_instances}) def apply_instance_deleted(event: InstanceDeleted, state: State) -> State: new_instances: Mapping[InstanceId, Instance] = { iid: inst for iid, inst in state.instances.items() if iid != event.instance_id } return state.model_copy(update={"instances": new_instances}) def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State: if isinstance(event.runner_status, RunnerShutdown): new_runners: Mapping[RunnerId, RunnerStatus] = { rid: rs for rid, rs in state.runners.items() if rid != event.runner_id } return state.model_copy(update={"runners": new_runners}) new_runners = { **state.runners, event.runner_id: event.runner_status, } return state.model_copy(update={"runners": new_runners}) def apply_node_timed_out(event: NodeTimedOut, state: State) -> State: topology = copy.deepcopy(state.topology) topology.remove_node(event.node_id) last_seen = { key: value for key, value in state.last_seen.items() if key != event.node_id } downloads = { key: value for key, value in state.downloads.items() if key != event.node_id } # Clean up all granular node mappings node_memory = { key: value for key, value in state.node_memory.items() if key != event.node_id } node_disk = { key: value for key, value in state.node_disk.items() if key != event.node_id } node_system = { key: value for key, value in state.node_system.items() if key != event.node_id } node_network = { key: value for key, value in state.node_network.items() if key != event.node_id } node_thunderbolt = { key: value for key, value in state.node_thunderbolt.items() if key != event.node_id } node_thunderbolt_bridge = { key: value for key, value in state.node_thunderbolt_bridge.items() if key != event.node_id } node_rdma_ctl = { key: value for key, value in state.node_rdma_ctl.items() if key != event.node_id } # Only recompute cycles if the leaving node had TB bridge enabled leaving_node_status = state.node_thunderbolt_bridge.get(event.node_id) leaving_node_had_tb_enabled = ( leaving_node_status is not None and leaving_node_status.enabled ) thunderbolt_bridge_cycles = ( topology.get_thunderbolt_bridge_cycles(node_thunderbolt_bridge, node_network) if leaving_node_had_tb_enabled else [list(cycle) for cycle in state.thunderbolt_bridge_cycles] ) return state.model_copy( update={ "downloads": downloads, "topology": topology, "last_seen": last_seen, "node_memory": node_memory, "node_disk": node_disk, "node_system": node_system, "node_network": node_network, "node_thunderbolt": node_thunderbolt, "node_thunderbolt_bridge": node_thunderbolt_bridge, "node_rdma_ctl": node_rdma_ctl, "thunderbolt_bridge_cycles": thunderbolt_bridge_cycles, } ) def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State: topology = copy.deepcopy(state.topology) topology.add_node(event.node_id) info = event.info # Build update dict with only the mappings that change update: dict[str, object] = { "last_seen": { **state.last_seen, event.node_id: datetime.fromisoformat(event.when), }, "topology": topology, } match info: case MacmonMetrics(): update["node_system"] = { **state.node_system, event.node_id: info.system_profile, } update["node_memory"] = {**state.node_memory, event.node_id: info.memory} case MemoryUsage(): update["node_memory"] = {**state.node_memory, event.node_id: info} case NodeDiskUsage(): update["node_disk"] = {**state.node_disk, event.node_id: info.disk_usage} case NodeConfig(): pass case MiscData(): current_identity = state.node_identities.get(event.node_id, NodeIdentity()) new_identity = current_identity.model_copy( update={"friendly_name": info.friendly_name} ) update["node_identities"] = { **state.node_identities, event.node_id: new_identity, } case StaticNodeInformation(): current_identity = state.node_identities.get(event.node_id, NodeIdentity()) new_identity = current_identity.model_copy( update={ "model_id": info.model, "chip_id": info.chip, "os_version": info.os_version, "os_build_version": info.os_build_version, } ) update["node_identities"] = { **state.node_identities, event.node_id: new_identity, } case NodeNetworkInterfaces(): update["node_network"] = { **state.node_network, event.node_id: NodeNetworkInfo(interfaces=info.ifaces), } case MacThunderboltIdentifiers(): update["node_thunderbolt"] = { **state.node_thunderbolt, event.node_id: NodeThunderboltInfo(interfaces=info.idents), } case MacThunderboltConnections(): conn_map = { tb_ident.domain_uuid: (nid, tb_ident.rdma_interface) for nid in state.node_thunderbolt for tb_ident in state.node_thunderbolt[nid].interfaces } as_rdma_conns = [ Connection( source=event.node_id, sink=conn_map[tb_conn.sink_uuid][0], edge=RDMAConnection( source_rdma_iface=conn_map[tb_conn.source_uuid][1], sink_rdma_iface=conn_map[tb_conn.sink_uuid][1], ), ) for tb_conn in info.conns if tb_conn.source_uuid in conn_map if tb_conn.sink_uuid in conn_map ] topology.replace_all_out_rdma_connections(event.node_id, as_rdma_conns) case ThunderboltBridgeInfo(): new_tb_bridge: dict[NodeId, ThunderboltBridgeStatus] = { **state.node_thunderbolt_bridge, event.node_id: info.status, } update["node_thunderbolt_bridge"] = new_tb_bridge # Only recompute cycles if the enabled status changed old_status = state.node_thunderbolt_bridge.get(event.node_id) old_enabled = old_status.enabled if old_status else False new_enabled = info.status.enabled if old_enabled != new_enabled: update["thunderbolt_bridge_cycles"] = ( topology.get_thunderbolt_bridge_cycles( new_tb_bridge, state.node_network ) ) case RdmaCtlStatus(): update["node_rdma_ctl"] = { **state.node_rdma_ctl, event.node_id: NodeRdmaCtlStatus(enabled=info.enabled), } return state.model_copy(update=update) def apply_topology_edge_created(event: TopologyEdgeCreated, state: State) -> State: topology = copy.deepcopy(state.topology) topology.add_connection(event.conn) return state.model_copy(update={"topology": topology}) def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> State: topology = copy.deepcopy(state.topology) topology.remove_connection(event.conn) # TODO: Clean up removing the reverse connection return state.model_copy(update={"topology": topology}) ================================================ FILE: src/exo/shared/constants.py ================================================ import os import sys from pathlib import Path from exo.utils.dashboard_path import find_dashboard, find_resources _EXO_HOME_ENV = os.environ.get("EXO_HOME", None) def _get_xdg_dir(env_var: str, fallback: str) -> Path: """Get XDG directory, prioritising EXO_HOME environment variable if its set. On non-Linux platforms, default to ~/.exo.""" if _EXO_HOME_ENV is not None: return Path.home() / _EXO_HOME_ENV if sys.platform != "linux": return Path.home() / ".exo" xdg_value = os.environ.get(env_var, None) if xdg_value is not None: return Path(xdg_value) / "exo" return Path.home() / fallback / "exo" EXO_CONFIG_HOME = _get_xdg_dir("XDG_CONFIG_HOME", ".config") EXO_DATA_HOME = _get_xdg_dir("XDG_DATA_HOME", ".local/share") EXO_CACHE_HOME = _get_xdg_dir("XDG_CACHE_HOME", ".cache") # Models directory (data) _EXO_MODELS_DIR_ENV = os.environ.get("EXO_MODELS_DIR", None) EXO_MODELS_DIR = ( EXO_DATA_HOME / "models" if _EXO_MODELS_DIR_ENV is None else Path.home() / _EXO_MODELS_DIR_ENV ) # Read-only search path for pre-downloaded models (colon-separated directories) _EXO_MODELS_PATH_ENV = os.environ.get("EXO_MODELS_PATH", None) EXO_MODELS_PATH: tuple[Path, ...] | None = ( tuple(Path(p).expanduser() for p in _EXO_MODELS_PATH_ENV.split(":") if p) if _EXO_MODELS_PATH_ENV is not None else None ) _RESOURCES_DIR_ENV = os.environ.get("EXO_RESOURCES_DIR", None) RESOURCES_DIR = ( find_resources() if _RESOURCES_DIR_ENV is None else Path.home() / _RESOURCES_DIR_ENV ) _DASHBOARD_DIR_ENV = os.environ.get("EXO_DASHBOARD_DIR", None) DASHBOARD_DIR = ( find_dashboard() if _DASHBOARD_DIR_ENV is None else Path.home() / _DASHBOARD_DIR_ENV ) # Log files (data/logs or cache) EXO_LOG_DIR = EXO_CACHE_HOME / "exo_log" EXO_LOG = EXO_LOG_DIR / "exo.log" EXO_TEST_LOG = EXO_CACHE_HOME / "exo_test.log" # Identity (config) EXO_NODE_ID_KEYPAIR = EXO_CONFIG_HOME / "node_id.keypair" EXO_CONFIG_FILE = EXO_CONFIG_HOME / "config.toml" # libp2p topics for event forwarding LIBP2P_LOCAL_EVENTS_TOPIC = "worker_events" LIBP2P_GLOBAL_EVENTS_TOPIC = "global_events" LIBP2P_ELECTION_MESSAGES_TOPIC = "election_message" LIBP2P_COMMANDS_TOPIC = "commands" EXO_MAX_CHUNK_SIZE = 512 * 1024 EXO_CUSTOM_MODEL_CARDS_DIR = EXO_DATA_HOME / "custom_model_cards" EXO_EVENT_LOG_DIR = EXO_DATA_HOME / "event_log" EXO_IMAGE_CACHE_DIR = EXO_CACHE_HOME / "images" EXO_TRACING_CACHE_DIR = EXO_CACHE_HOME / "traces" EXO_ENABLE_IMAGE_MODELS = ( os.getenv("EXO_ENABLE_IMAGE_MODELS", "false").lower() == "true" ) EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true" EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true" EXO_MAX_CONCURRENT_REQUESTS = int(os.getenv("EXO_MAX_CONCURRENT_REQUESTS", "8")) ================================================ FILE: src/exo/shared/election.py ================================================ from typing import Self import anyio from anyio import ( CancelScope, Event, get_cancelled_exc_class, ) from loguru import logger from exo.routing.connection_message import ConnectionMessage from exo.shared.types.commands import ForwarderCommand from exo.shared.types.common import NodeId, SessionId from exo.utils.channels import Receiver, Sender from exo.utils.pydantic_ext import CamelCaseModel from exo.utils.task_group import TaskGroup DEFAULT_ELECTION_TIMEOUT = 3.0 class ElectionMessage(CamelCaseModel): clock: int seniority: int proposed_session: SessionId commands_seen: int # Could eventually include a list of neighbour nodes for centrality def __lt__(self, other: Self) -> bool: if self.clock != other.clock: return self.clock < other.clock if self.seniority != other.seniority: return self.seniority < other.seniority elif self.commands_seen != other.commands_seen: return self.commands_seen < other.commands_seen else: return ( self.proposed_session.master_node_id < other.proposed_session.master_node_id ) class ElectionResult(CamelCaseModel): session_id: SessionId won_clock: int is_new_master: bool class Election: def __init__( self, node_id: NodeId, *, election_message_receiver: Receiver[ElectionMessage], election_message_sender: Sender[ElectionMessage], election_result_sender: Sender[ElectionResult], connection_message_receiver: Receiver[ConnectionMessage], command_receiver: Receiver[ForwarderCommand], is_candidate: bool = True, seniority: int = 0, ): # If we aren't a candidate, simply don't increment seniority. # For reference: This node can be elected master if all nodes are not master candidates # Any master candidate will automatically win out over this node. self.seniority = seniority if is_candidate else -1 self.clock = 0 self.node_id = node_id self.commands_seen = 0 # Every node spawns as master self.current_session: SessionId = SessionId( master_node_id=node_id, election_clock=0 ) # Senders/Receivers self._em_sender = election_message_sender self._em_receiver = election_message_receiver self._er_sender = election_result_sender self._cm_receiver = connection_message_receiver self._co_receiver = command_receiver # Campaign state self._candidates: list[ElectionMessage] = [] self._campaign_cancel_scope: CancelScope | None = None self._campaign_done: Event | None = None self._tg = TaskGroup() async def run(self): logger.info("Starting Election") try: async with self._tg as tg: tg.start_soon(self._election_receiver) tg.start_soon(self._connection_receiver) tg.start_soon(self._command_counter) # And start an election immediately, that instantly resolves candidates: list[ElectionMessage] = [] logger.debug("Starting initial campaign") self._candidates = candidates await self._campaign(candidates, campaign_timeout=0.0) logger.debug("Initial campaign finished") finally: # Cancel and wait for the last election to end if self._campaign_cancel_scope is not None: logger.debug("Cancelling campaign") self._campaign_cancel_scope.cancel() if self._campaign_done is not None: logger.debug("Waiting for campaign to finish") await self._campaign_done.wait() logger.debug("Campaign cancelled and finished") logger.info("Election shutdown") async def elect(self, em: ElectionMessage) -> None: logger.debug(f"Electing: {em}") is_new_master = em.proposed_session != self.current_session self.current_session = em.proposed_session logger.debug(f"Current session: {self.current_session}") await self._er_sender.send( ElectionResult( won_clock=em.clock, session_id=em.proposed_session, is_new_master=is_new_master, ) ) async def shutdown(self) -> None: self._tg.cancel_tasks() async def _election_receiver(self) -> None: with self._em_receiver as election_messages: async for message in election_messages: logger.debug(f"Election message received: {message}") if message.proposed_session.master_node_id == self.node_id: logger.debug("Dropping message from ourselves") # Drop messages from us (See exo.routing.router) continue # If a new round is starting, we participate if message.clock > self.clock: self.clock = message.clock logger.debug(f"New clock: {self.clock}") logger.debug("Starting new campaign") candidates: list[ElectionMessage] = [message] logger.debug(f"Candidates: {candidates}") logger.debug(f"Current candidates: {self._candidates}") self._candidates = candidates logger.debug(f"New candidates: {self._candidates}") logger.debug("Starting new campaign") self._tg.start_soon( self._campaign, candidates, DEFAULT_ELECTION_TIMEOUT ) logger.debug("Campaign started") continue # Dismiss old messages if message.clock < self.clock: logger.debug(f"Dropping old message: {message}") continue logger.debug(f"Election added candidate {message}") # Now we are processing this rounds messages - including the message that triggered this round. self._candidates.append(message) async def _connection_receiver(self) -> None: with self._cm_receiver as connection_messages: async for first in connection_messages: # Delay after connection message for time to symmetrically setup await anyio.sleep(0.2) rest = connection_messages.collect() logger.debug( f"Connection messages received: {first} followed by {rest}" ) logger.debug(f"Current clock: {self.clock}") # These messages are strictly peer to peer self.clock += 1 logger.debug(f"New clock: {self.clock}") candidates: list[ElectionMessage] = [] self._candidates = candidates logger.debug("Starting new campaign") self._tg.start_soon( self._campaign, candidates, DEFAULT_ELECTION_TIMEOUT ) logger.debug("Campaign started") logger.debug("Connection message added") async def _command_counter(self) -> None: with self._co_receiver as commands: async for _command in commands: self.commands_seen += 1 async def _campaign( self, candidates: list[ElectionMessage], campaign_timeout: float ) -> None: clock = self.clock # Kill the old campaign if self._campaign_cancel_scope: logger.info("Cancelling other campaign") self._campaign_cancel_scope.cancel() if self._campaign_done: logger.info("Waiting for other campaign to finish") await self._campaign_done.wait() done = Event() self._campaign_done = done scope = CancelScope() self._campaign_cancel_scope = scope try: with scope: logger.debug(f"Election {clock} started") status = self._election_status(clock) candidates.append(status) await self._em_sender.send(status) logger.debug(f"Sleeping for {campaign_timeout} seconds") await anyio.sleep(campaign_timeout) # minor hack - rebroadcast status in case anyone has missed it. await self._em_sender.send(status) logger.debug("Woke up from sleep") # add an anyio checkpoint - anyio.lowlevel.chekpoint() or checkpoint_if_cancelled() is preferred, but wasn't typechecking last I checked await anyio.sleep(0) # Election finished! elected = max(candidates) logger.debug(f"Election queue {candidates}") logger.debug(f"Elected: {elected}") if ( self.node_id == elected.proposed_session.master_node_id and self.seniority >= 0 ): logger.debug( f"Node is a candidate and seniority is {self.seniority}" ) self.seniority = max(self.seniority, len(candidates)) logger.debug(f"New seniority: {self.seniority}") else: logger.debug( f"Node is not a candidate or seniority is not {self.seniority}" ) logger.debug( f"Election finished, new SessionId({elected.proposed_session}) with queue {candidates}" ) logger.debug("Sending election result") await self.elect(elected) logger.debug("Election result sent") except get_cancelled_exc_class(): logger.debug(f"Election {clock} cancelled") finally: logger.debug(f"Election {clock} finally") if self._campaign_cancel_scope is scope: self._campaign_cancel_scope = None logger.debug("Setting done event") done.set() logger.debug("Done event set") def _election_status(self, clock: int | None = None) -> ElectionMessage: c = self.clock if clock is None else clock return ElectionMessage( proposed_session=( self.current_session if self.current_session.master_node_id == self.node_id else SessionId(master_node_id=self.node_id, election_clock=c) ), clock=c, seniority=self.seniority, commands_seen=self.commands_seen, ) ================================================ FILE: src/exo/shared/logging.py ================================================ import logging import sys from collections.abc import Iterator from pathlib import Path import zstandard from hypercorn import Config from hypercorn.logging import Logger as HypercornLogger from loguru import logger _MAX_LOG_ARCHIVES = 5 def _zstd_compress(filepath: str) -> None: source = Path(filepath) dest = source.with_suffix(source.suffix + ".zst") cctx = zstandard.ZstdCompressor() with open(source, "rb") as f_in, open(dest, "wb") as f_out: cctx.copy_stream(f_in, f_out) source.unlink() def _once_then_never() -> Iterator[bool]: yield True while True: yield False class InterceptLogger(HypercornLogger): def __init__(self, config: Config): super().__init__(config) assert self.error_logger self.error_logger.handlers = [_InterceptHandler()] class _InterceptHandler(logging.Handler): def emit(self, record: logging.LogRecord): try: level = logger.level(record.levelname).name except ValueError: level = record.levelno logger.opt(depth=3, exception=record.exc_info).log(level, record.getMessage()) def logger_setup(log_file: Path | None, verbosity: int = 0): """Set up logging for this process - formatting, file handles, verbosity and output""" logging.getLogger("exo_pyo3_bindings").setLevel(logging.WARNING) logging.getLogger("httpx").setLevel(logging.WARNING) logging.getLogger("httpcore").setLevel(logging.WARNING) logger.remove() # replace all stdlib loggers with _InterceptHandlers that log to loguru logging.basicConfig(handlers=[_InterceptHandler()], level=0) if verbosity == 0: logger.add( sys.__stderr__, # type: ignore format="[ {time:hh:mm:ss.SSSSA} | {level: <8}] {message}", level="INFO", colorize=True, enqueue=True, ) else: logger.add( sys.__stderr__, # type: ignore format="[ {time:HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} ] {message}", level="DEBUG", colorize=True, enqueue=True, ) if log_file: rotate_once = _once_then_never() logger.add( log_file, format="[ {time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} ] {message}", level="INFO", colorize=False, enqueue=True, rotation=lambda _, __: next(rotate_once), retention=_MAX_LOG_ARCHIVES, compression=_zstd_compress, ) def logger_cleanup(): """Flush all queues before shutting down so any in-flight logs are written to disk""" logger.complete() """ --- TODO: Capture MLX Log output: import contextlib import sys from loguru import logger class StreamToLogger: def __init__(self, level="INFO"): self._level = level def write(self, buffer): for line in buffer.rstrip().splitlines(): logger.opt(depth=1).log(self._level, line.rstrip()) def flush(self): pass logger.remove() logger.add(sys.__stdout__) stream = StreamToLogger() with contextlib.redirect_stdout(stream): print("Standard output is sent to added handlers.") """ ================================================ FILE: src/exo/shared/models/model_cards.py ================================================ from enum import Enum from typing import Annotated, Any import aiofiles import aiofiles.os as aios import tomlkit from anyio import Path, open_file from huggingface_hub import model_info from loguru import logger from pydantic import ( AliasChoices, BaseModel, Field, PositiveInt, ValidationError, field_validator, model_validator, ) from tomlkit.exceptions import TOMLKitError from exo.shared.constants import ( EXO_CUSTOM_MODEL_CARDS_DIR, EXO_ENABLE_IMAGE_MODELS, RESOURCES_DIR, ) from exo.shared.types.common import ModelId from exo.shared.types.memory import Memory from exo.utils.pydantic_ext import CamelCaseModel # kinda ugly... # TODO: load search path from config.toml _custom_cards_dir = Path(str(EXO_CUSTOM_MODEL_CARDS_DIR)) CARD_SEARCH_PATH = [ Path(RESOURCES_DIR) / "inference_model_cards", Path(RESOURCES_DIR) / "image_model_cards", _custom_cards_dir, ] _card_cache: dict[ModelId, "ModelCard"] = {} async def _refresh_card_cache(): for path in CARD_SEARCH_PATH: async for toml_file in path.rglob("*.toml"): try: card = await ModelCard.load_from_path(toml_file) if card.model_id not in _card_cache: _card_cache[card.model_id] = card except (ValidationError, TOMLKitError): pass def _is_image_card(card: "ModelCard") -> bool: return any(t in (ModelTask.TextToImage, ModelTask.ImageToImage) for t in card.tasks) async def get_model_cards() -> list["ModelCard"]: if len(_card_cache) == 0: await _refresh_card_cache() if EXO_ENABLE_IMAGE_MODELS: return list(_card_cache.values()) return [c for c in _card_cache.values() if not _is_image_card(c)] class ModelTask(str, Enum): TextGeneration = "TextGeneration" TextToImage = "TextToImage" ImageToImage = "ImageToImage" class ComponentInfo(CamelCaseModel): component_name: str component_path: str storage_size: Memory n_layers: PositiveInt | None = None can_shard: bool safetensors_index_filename: str | None = None class ModelCard(CamelCaseModel): model_id: ModelId storage_size: Memory n_layers: PositiveInt hidden_size: PositiveInt supports_tensor: bool num_key_value_heads: PositiveInt | None = None tasks: list[ModelTask] components: list[ComponentInfo] | None = None family: str = "" quantization: str = "" base_model: str = "" capabilities: list[str] = [] uses_cfg: bool = False trust_remote_code: bool = True @field_validator("tasks", mode="before") @classmethod def _validate_tasks(cls, v: list[str | ModelTask]) -> list[ModelTask]: return [item if isinstance(item, ModelTask) else ModelTask(item) for item in v] async def save(self, path: Path) -> None: async with await open_file(path, "w") as f: py = self.model_dump(exclude_none=True) data = tomlkit.dumps(py) # pyright: ignore[reportUnknownMemberType] await f.write(data) async def save_to_custom_dir(self) -> None: await aios.makedirs(str(_custom_cards_dir), exist_ok=True) await self.save(_custom_cards_dir / (self.model_id.normalize() + ".toml")) @staticmethod async def load_from_path(path: Path) -> "ModelCard": async with await open_file(path, "r") as f: py = tomlkit.loads(await f.read()) return ModelCard.model_validate(py) # Is it okay that model card.load defaults to network access if the card doesn't exist? do we want to be more explicit here? @staticmethod async def load(model_id: ModelId) -> "ModelCard": if model_id not in _card_cache: await _refresh_card_cache() if (mc := _card_cache.get(model_id)) is not None: return mc return await ModelCard.fetch_from_hf(model_id) @staticmethod async def fetch_from_hf(model_id: ModelId) -> "ModelCard": """Fetches storage size and number of layers for a Hugging Face model, returns Pydantic ModelMeta.""" # TODO: failure if files do not exist config_data = await fetch_config_data(model_id) num_layers = config_data.layer_count mem_size_bytes = await fetch_safetensors_size(model_id) mc = ModelCard( model_id=ModelId(model_id), storage_size=mem_size_bytes, n_layers=num_layers, hidden_size=config_data.hidden_size or 0, supports_tensor=config_data.supports_tensor, num_key_value_heads=config_data.num_key_value_heads, tasks=[ModelTask.TextGeneration], trust_remote_code=False, ) await mc.save_to_custom_dir() _card_cache[model_id] = mc return mc async def delete_custom_card(model_id: ModelId) -> bool: """Delete a user-added custom model card. Returns True if deleted.""" card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml") if await card_path.exists(): await card_path.unlink() _card_cache.pop(model_id, None) return True return False def is_custom_card(model_id: ModelId) -> bool: """Check if a model card exists in the custom cards directory.""" import os card_path = Path(str(EXO_CUSTOM_MODEL_CARDS_DIR)) / ( ModelId(model_id).normalize() + ".toml" ) return os.path.isfile(str(card_path)) class ConfigData(BaseModel): model_config = {"extra": "ignore"} # Allow unknown fields architectures: list[str] | None = None hidden_size: Annotated[int, Field(ge=0)] | None = None num_key_value_heads: PositiveInt | None = None layer_count: int = Field( validation_alias=AliasChoices( "num_hidden_layers", "num_layers", "n_layer", "n_layers", "num_decoder_layers", "decoder_layers", ) ) @property def supports_tensor(self) -> bool: return self.architectures in [ ["Glm4MoeLiteForCausalLM"], ["GlmMoeDsaForCausalLM"], ["DeepseekV32ForCausalLM"], ["DeepseekV3ForCausalLM"], ["Qwen3NextForCausalLM"], ["Qwen3MoeForCausalLM"], ["Qwen3_5MoeForConditionalGeneration"], ["Qwen3_5ForConditionalGeneration"], ["MiniMaxM2ForCausalLM"], ["LlamaForCausalLM"], ["GptOssForCausalLM"], ["Step3p5ForCausalLM"], ["NemotronHForCausalLM"], ] @model_validator(mode="before") @classmethod def defer_to_text_config(cls, data: dict[str, Any]): text_config = data.get("text_config") if text_config is None: return data for field in [ "architectures", "hidden_size", "num_key_value_heads", "num_hidden_layers", "num_layers", "n_layer", "n_layers", "num_decoder_layers", "decoder_layers", ]: if (val := text_config.get(field)) is not None: # pyright: ignore[reportAny] data[field] = val return data async def fetch_config_data(model_id: ModelId) -> ConfigData: """Downloads and parses config.json for a model.""" from exo.download.download_utils import ( download_file_with_retry, ensure_models_dir, ) target_dir = (await ensure_models_dir()) / model_id.normalize() await aios.makedirs(target_dir, exist_ok=True) config_path = await download_file_with_retry( model_id, "main", "config.json", target_dir, lambda curr_bytes, total_bytes, is_renamed: logger.debug( f"Downloading config.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})" ), ) async with aiofiles.open(config_path, "r") as f: return ConfigData.model_validate_json(await f.read()) async def fetch_safetensors_size(model_id: ModelId) -> Memory: """Gets model size from safetensors index or falls back to HF API.""" from exo.download.download_utils import ( download_file_with_retry, ensure_models_dir, ) from exo.shared.types.worker.downloads import ModelSafetensorsIndex target_dir = (await ensure_models_dir()) / model_id.normalize() await aios.makedirs(target_dir, exist_ok=True) index_path = await download_file_with_retry( model_id, "main", "model.safetensors.index.json", target_dir, lambda curr_bytes, total_bytes, is_renamed: logger.debug( f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})" ), ) async with aiofiles.open(index_path, "r") as f: index_data = ModelSafetensorsIndex.model_validate_json(await f.read()) metadata = index_data.metadata if metadata is not None: return Memory.from_bytes(metadata.total_size) info = model_info(model_id) if info.safetensors is None: raise ValueError(f"No safetensors info found for {model_id}") return Memory.from_bytes(info.safetensors.total) ================================================ FILE: src/exo/shared/tests/__init__.py ================================================ # Test package for shared utilities ================================================ FILE: src/exo/shared/tests/conftest.py ================================================ """Pytest configuration and shared fixtures for shared package tests.""" import asyncio from typing import Generator import pytest from _pytest.logging import LogCaptureFixture from loguru import logger from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask from exo.shared.types.memory import Memory from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata @pytest.fixture(scope="session") def event_loop() -> Generator[asyncio.AbstractEventLoop, None, None]: """Create an event loop for the test session.""" loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) yield loop loop.close() @pytest.fixture(autouse=True) def reset_event_loop(): """Reset the event loop for each test to ensure clean state.""" # This ensures each test gets a fresh event loop state def get_pipeline_shard_metadata( model_id: ModelId, device_rank: int, world_size: int = 1 ) -> ShardMetadata: return PipelineShardMetadata( model_card=ModelCard( model_id=model_id, storage_size=Memory.from_mb(100000), n_layers=32, hidden_size=1000, supports_tensor=True, tasks=[ModelTask.TextGeneration], ), device_rank=device_rank, world_size=world_size, start_layer=0, end_layer=32, n_layers=32, ) @pytest.fixture def caplog(caplog: LogCaptureFixture): handler_id = logger.add( caplog.handler, format="{message}", level=0, filter=lambda record: record["level"].no >= caplog.handler.level, enqueue=True, # Set to 'True' if your test is spawning child processes. ) yield caplog logger.remove(handler_id) ================================================ FILE: src/exo/shared/tests/test_apply/test_apply_node_download.py ================================================ from exo.shared.apply import apply_node_download_progress from exo.shared.tests.conftest import get_pipeline_shard_metadata from exo.shared.types.common import NodeId from exo.shared.types.events import NodeDownloadProgress from exo.shared.types.memory import Memory from exo.shared.types.state import State from exo.shared.types.worker.downloads import DownloadCompleted from exo.worker.tests.constants import MODEL_A_ID, MODEL_B_ID def test_apply_node_download_progress(): state = State() shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) event = DownloadCompleted( node_id=NodeId("node-1"), shard_metadata=shard1, total=Memory(), ) new_state = apply_node_download_progress( NodeDownloadProgress(download_progress=event), state ) assert new_state.downloads == {NodeId("node-1"): [event]} def test_apply_two_node_download_progress(): shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard2 = get_pipeline_shard_metadata(MODEL_B_ID, device_rank=0, world_size=2) event1 = DownloadCompleted( node_id=NodeId("node-1"), shard_metadata=shard1, total=Memory(), ) event2 = DownloadCompleted( node_id=NodeId("node-1"), shard_metadata=shard2, total=Memory(), ) state = State(downloads={NodeId("node-1"): [event1]}) new_state = apply_node_download_progress( NodeDownloadProgress(download_progress=event2), state ) assert new_state.downloads == {NodeId("node-1"): [event1, event2]} ================================================ FILE: src/exo/shared/tests/test_apply/test_apply_runner_deleted.py ================================================ from exo.shared.apply import apply_runner_status_updated from exo.shared.types.events import RunnerStatusUpdated from exo.shared.types.state import State from exo.shared.types.worker.runners import RunnerId, RunnerIdle, RunnerShutdown def test_apply_runner_shutdown_removes_runner(): runner_id = RunnerId() state = State(runners={runner_id: RunnerIdle()}) new_state = apply_runner_status_updated( RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerShutdown()), state ) assert runner_id not in new_state.runners def test_apply_runner_status_updated_adds_runner(): runner_id = RunnerId() state = State() new_state = apply_runner_status_updated( RunnerStatusUpdated(runner_id=runner_id, runner_status=RunnerIdle()), state ) assert runner_id in new_state.runners ================================================ FILE: src/exo/shared/tests/test_election.py ================================================ import pytest from anyio import create_task_group, fail_after, move_on_after from exo.routing.connection_message import ConnectionMessage from exo.shared.election import Election, ElectionMessage, ElectionResult from exo.shared.types.commands import ForwarderCommand, TestCommand from exo.shared.types.common import NodeId, SessionId, SystemId from exo.utils.channels import channel # ======= # # Helpers # # ======= # def em( clock: int, seniority: int, node_id: str, commands_seen: int = 0, election_clock: int | None = None, ) -> ElectionMessage: """ Helper to build ElectionMessages for a given proposer node. The new API carries a proposed SessionId (master_node_id + election_clock). By default we use the same value for election_clock as the 'clock' of the round. """ return ElectionMessage( clock=clock, seniority=seniority, proposed_session=SessionId( master_node_id=NodeId(node_id), election_clock=clock if election_clock is None else election_clock, ), commands_seen=commands_seen, ) # ======================================= # # TESTS # # ======================================= # @pytest.fixture(autouse=True) def fast_election_timeout(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr("exo.shared.election.DEFAULT_ELECTION_TIMEOUT", 0.1) @pytest.mark.anyio async def test_single_round_broadcasts_and_updates_seniority_on_self_win() -> None: """ Start a round by injecting an ElectionMessage with higher clock. With only our node effectively 'winning', we should broadcast once and update seniority. """ # Outbound election messages from the Election (we'll observe these) em_out_tx, em_out_rx = channel[ElectionMessage]() # Inbound election messages to the Election (we'll inject these) em_in_tx, em_in_rx = channel[ElectionMessage]() # Election results produced by the Election (we'll observe these) er_tx, er_rx = channel[ElectionResult]() # Connection messages cm_tx, cm_rx = channel[ConnectionMessage]() # Commands co_tx, co_rx = channel[ForwarderCommand]() election = Election( node_id=NodeId("B"), election_message_receiver=em_in_rx, election_message_sender=em_out_tx, election_result_sender=er_tx, connection_message_receiver=cm_rx, command_receiver=co_rx, is_candidate=True, ) async with create_task_group() as tg: with fail_after(2): tg.start_soon(election.run) # Trigger new round at clock=1 (peer announces it) await em_in_tx.send(em(clock=1, seniority=0, node_id="A")) # Expect our broadcast back to the peer side for this round only while True: got = await em_out_rx.receive() if got.clock == 1 and got.proposed_session.master_node_id == NodeId( "B" ): break # Wait for the round to finish and produce an ElectionResult result = await er_rx.receive() assert result.session_id.master_node_id == NodeId("B") # We spawned as master; electing ourselves again is not "new master". assert result.is_new_master is False # Close inbound streams to end the receivers (and run()) em_in_tx.close() cm_tx.close() co_tx.close() # We should have updated seniority to 2 (A + B). assert election.seniority == 2 @pytest.mark.anyio async def test_peer_with_higher_seniority_wins_and_we_switch_master() -> None: """ If a peer with clearly higher seniority participates in the round, they should win. We should broadcast our status exactly once for this round, then switch master. """ em_out_tx, em_out_rx = channel[ElectionMessage]() em_in_tx, em_in_rx = channel[ElectionMessage]() er_tx, er_rx = channel[ElectionResult]() cm_tx, cm_rx = channel[ConnectionMessage]() co_tx, co_rx = channel[ForwarderCommand]() election = Election( node_id=NodeId("ME"), election_message_receiver=em_in_rx, election_message_sender=em_out_tx, election_result_sender=er_tx, connection_message_receiver=cm_rx, command_receiver=co_rx, is_candidate=True, ) async with create_task_group() as tg: with fail_after(2): tg.start_soon(election.run) # Start round with peer's message (higher seniority) await em_in_tx.send(em(clock=1, seniority=10, node_id="PEER")) # We should still broadcast our status exactly once for this round while True: got = await em_out_rx.receive() if got.clock == 1: assert got.seniority == 0 break # After the timeout, election result for clock=1 should report the peer as master # (Skip any earlier result from the boot campaign at clock=0 by filtering on election_clock) while True: result = await er_rx.receive() if result.session_id.election_clock == 1: break assert result.session_id.master_node_id == NodeId("PEER") assert result.is_new_master is True em_in_tx.close() cm_tx.close() co_tx.close() # We lost → seniority unchanged assert election.seniority == 0 @pytest.mark.anyio async def test_ignores_older_messages() -> None: """ Messages with a lower clock than the current round are ignored by the receiver. Expect exactly one broadcast for the higher clock round. """ em_out_tx, em_out_rx = channel[ElectionMessage]() em_in_tx, em_in_rx = channel[ElectionMessage]() er_tx, _er_rx = channel[ElectionResult]() cm_tx, cm_rx = channel[ConnectionMessage]() co_tx, co_rx = channel[ForwarderCommand]() election = Election( node_id=NodeId("ME"), election_message_receiver=em_in_rx, election_message_sender=em_out_tx, election_result_sender=er_tx, connection_message_receiver=cm_rx, command_receiver=co_rx, is_candidate=True, ) async with create_task_group() as tg: with fail_after(2): tg.start_soon(election.run) # Newer round arrives first -> triggers campaign at clock=2 await em_in_tx.send(em(clock=2, seniority=0, node_id="A")) while True: first = await em_out_rx.receive() if first.clock == 2: break # Older message (clock=1) must be ignored (no second broadcast) await em_in_tx.send(em(clock=1, seniority=999, node_id="B")) got_second = False with move_on_after(0.05): _ = await em_out_rx.receive() got_second = True assert not got_second, "Should not receive a broadcast for an older round" em_in_tx.close() cm_tx.close() co_tx.close() # Not asserting on the result; focus is on ignore behavior. @pytest.mark.anyio async def test_two_rounds_emit_two_broadcasts_and_increment_clock() -> None: """ Two successive rounds → two broadcasts. Second round triggered by a higher-clock message. """ em_out_tx, em_out_rx = channel[ElectionMessage]() em_in_tx, em_in_rx = channel[ElectionMessage]() er_tx, _er_rx = channel[ElectionResult]() cm_tx, cm_rx = channel[ConnectionMessage]() co_tx, co_rx = channel[ForwarderCommand]() election = Election( node_id=NodeId("ME"), election_message_receiver=em_in_rx, election_message_sender=em_out_tx, election_result_sender=er_tx, connection_message_receiver=cm_rx, command_receiver=co_rx, is_candidate=True, ) async with create_task_group() as tg: with fail_after(2): tg.start_soon(election.run) # Round 1 at clock=1 await em_in_tx.send(em(clock=1, seniority=0, node_id="X")) while True: m1 = await em_out_rx.receive() if m1.clock == 1: break # Round 2 at clock=2 await em_in_tx.send(em(clock=2, seniority=0, node_id="Y")) while True: m2 = await em_out_rx.receive() if m2.clock == 2: break em_in_tx.close() cm_tx.close() co_tx.close() # Not asserting on who won; just that both rounds were broadcast. @pytest.mark.anyio async def test_promotion_new_seniority_counts_participants() -> None: """ When we win against two peers in the same round, our seniority becomes max(existing, number_of_candidates). With existing=0: expect 3 (us + A + B). """ em_out_tx, em_out_rx = channel[ElectionMessage]() em_in_tx, em_in_rx = channel[ElectionMessage]() er_tx, er_rx = channel[ElectionResult]() cm_tx, cm_rx = channel[ConnectionMessage]() co_tx, co_rx = channel[ForwarderCommand]() election = Election( node_id=NodeId("ME"), election_message_receiver=em_in_rx, election_message_sender=em_out_tx, election_result_sender=er_tx, connection_message_receiver=cm_rx, command_receiver=co_rx, is_candidate=True, ) async with create_task_group() as tg: with fail_after(2): tg.start_soon(election.run) # Start round at clock=7 with two peer participants await em_in_tx.send(em(clock=7, seniority=0, node_id="A")) await em_in_tx.send(em(clock=7, seniority=0, node_id="B")) # We should see exactly one broadcast from us for this round while True: got = await em_out_rx.receive() if got.clock == 7 and got.proposed_session.master_node_id == NodeId( "ME" ): break # Wait for the election to finish so seniority updates _ = await er_rx.receive() em_in_tx.close() cm_tx.close() co_tx.close() # We + A + B = 3 → new seniority expected to be 3 assert election.seniority == 3 @pytest.mark.anyio async def test_connection_message_triggers_new_round_broadcast() -> None: """ A connection message increments the clock and starts a new campaign. We should observe a broadcast at the incremented clock. """ em_out_tx, em_out_rx = channel[ElectionMessage]() em_in_tx, em_in_rx = channel[ElectionMessage]() er_tx, _er_rx = channel[ElectionResult]() cm_tx, cm_rx = channel[ConnectionMessage]() co_tx, co_rx = channel[ForwarderCommand]() election = Election( node_id=NodeId("ME"), election_message_receiver=em_in_rx, election_message_sender=em_out_tx, election_result_sender=er_tx, connection_message_receiver=cm_rx, command_receiver=co_rx, is_candidate=True, ) async with create_task_group() as tg: with fail_after(2): tg.start_soon(election.run) # Send any connection message object; we close quickly to cancel before result creation await cm_tx.send(ConnectionMessage(node_id=NodeId(), connected=True)) # Expect a broadcast for the new round at clock=1 while True: got = await em_out_rx.receive() if got.clock == 1 and got.proposed_session.master_node_id == NodeId( "ME" ): break # Close promptly to avoid waiting for campaign completion em_in_tx.close() cm_tx.close() co_tx.close() # After cancellation (before election finishes), no seniority changes asserted here. @pytest.mark.anyio async def test_tie_breaker_prefers_node_with_more_commands_seen() -> None: """ With equal seniority, the node that has seen more commands should win the election. We increase our local 'commands_seen' by sending TestCommand()s before triggering the round. """ em_out_tx, em_out_rx = channel[ElectionMessage]() em_in_tx, em_in_rx = channel[ElectionMessage]() er_tx, er_rx = channel[ElectionResult]() cm_tx, cm_rx = channel[ConnectionMessage]() co_tx, co_rx = channel[ForwarderCommand]() me = NodeId("ME") election = Election( node_id=me, election_message_receiver=em_in_rx, election_message_sender=em_out_tx, election_result_sender=er_tx, connection_message_receiver=cm_rx, command_receiver=co_rx, is_candidate=True, seniority=0, ) async with create_task_group() as tg: with fail_after(2): tg.start_soon(election.run) # Pump local commands so our commands_seen is high before the round starts for _ in range(50): await co_tx.send( ForwarderCommand(origin=SystemId("SOMEONE"), command=TestCommand()) ) # Trigger a round at clock=1 with a peer of equal seniority but fewer commands await em_in_tx.send( em(clock=1, seniority=0, node_id="PEER", commands_seen=5) ) # Observe our broadcast for this round (to ensure we've joined the round) while True: got = await em_out_rx.receive() if got.clock == 1 and got.proposed_session.master_node_id == me: # We don't assert exact count, just that we've participated this round. break # The elected result for clock=1 should be us due to higher commands_seen while True: result = await er_rx.receive() if result.session_id.master_node_id == me: assert result.session_id.election_clock in (0, 1) break em_in_tx.close() cm_tx.close() co_tx.close() ================================================ FILE: src/exo/shared/tests/test_node_id_persistence.py ================================================ import contextlib import multiprocessing import os from multiprocessing import Event, Queue, Semaphore from multiprocessing.process import BaseProcess from multiprocessing.queues import Queue as QueueT from multiprocessing.synchronize import Event as EventT from multiprocessing.synchronize import Semaphore as SemaphoreT from loguru import logger from pytest import LogCaptureFixture, mark from exo.routing.router import get_node_id_keypair from exo.shared.constants import EXO_NODE_ID_KEYPAIR NUM_CONCURRENT_PROCS = 10 def _get_keypair_concurrent_subprocess_task( sem: SemaphoreT, ev: EventT, queue: QueueT[bytes] ) -> None: # synchronise with parent process sem.release() # wait to be told to begin simultaneous read ev.wait() queue.put(get_node_id_keypair().to_bytes()) def _get_keypair_concurrent(num_procs: int) -> bytes: assert num_procs > 0 sem = Semaphore(0) ev = Event() queue: QueueT[bytes] = Queue(maxsize=num_procs) # make parent process wait for all subprocesses to start logger.info(f"PARENT: Starting {num_procs} subprocesses") ps: list[BaseProcess] = [] for _ in range(num_procs): p = multiprocessing.get_context("fork").Process( target=_get_keypair_concurrent_subprocess_task, args=(sem, ev, queue) ) ps.append(p) p.start() for _ in range(num_procs): sem.acquire() # start all the sub processes simultaneously logger.info("PARENT: Beginning read") ev.set() # wait until all subprocesses are done & read results for p in ps: p.join() # check that the input/output order match, and that # all subprocesses end up reading the same file logger.info("PARENT: Checking consistency") keypair: bytes | None = None qsize = 0 # cannot use Queue.qsize due to MacOS incompatibility :( while not queue.empty(): qsize += 1 temp_keypair = queue.get() if keypair is None: keypair = temp_keypair else: assert keypair == temp_keypair assert num_procs == qsize return keypair # pyright: ignore[reportReturnType] def _delete_if_exists(p: str | bytes | os.PathLike[str] | os.PathLike[bytes]): with contextlib.suppress(OSError): os.remove(p) @mark.skip(reason="this functionality is currently disabled but may return in future") def test_node_id_fetching(caplog: LogCaptureFixture): reps = 10 # delete current file and write a new one _delete_if_exists(EXO_NODE_ID_KEYPAIR) kp = _get_keypair_concurrent(NUM_CONCURRENT_PROCS) with caplog.at_level(101): # supress logs # make sure that continuous fetches return the same value for _ in range(reps): assert kp == _get_keypair_concurrent(NUM_CONCURRENT_PROCS) # make sure that after deleting, we are not fetching the same value _delete_if_exists(EXO_NODE_ID_KEYPAIR) for _ in range(reps): assert kp != _get_keypair_concurrent(NUM_CONCURRENT_PROCS) ================================================ FILE: src/exo/shared/tests/test_resolve_reasoning_params.py ================================================ import pytest from exo.shared.types.text_generation import resolve_reasoning_params def test_both_none_returns_none_none() -> None: assert resolve_reasoning_params(None, None) == (None, None) def test_both_set_passes_through_unchanged() -> None: assert resolve_reasoning_params("high", True) == ("high", True) assert resolve_reasoning_params("none", True) == ("none", True) assert resolve_reasoning_params("low", False) == ("low", False) def test_enable_thinking_true_derives_medium() -> None: assert resolve_reasoning_params(None, True) == ("medium", True) def test_enable_thinking_false_derives_none() -> None: assert resolve_reasoning_params(None, False) == ("none", False) def test_reasoning_effort_none_derives_thinking_false() -> None: assert resolve_reasoning_params("none", None) == ("none", False) @pytest.mark.parametrize("effort", ["minimal", "low", "medium", "high", "xhigh"]) def test_non_none_effort_derives_thinking_true(effort: str) -> None: assert resolve_reasoning_params(effort, None) == (effort, True) # pyright: ignore[reportArgumentType] ================================================ FILE: src/exo/shared/tests/test_state_serialization.py ================================================ from exo.shared.types.common import NodeId from exo.shared.types.multiaddr import Multiaddr from exo.shared.types.state import State from exo.shared.types.topology import Connection, SocketConnection def test_state_serialization_roundtrip() -> None: """Verify that State → JSON → State round-trip preserves topology.""" # --- build a simple state ------------------------------------------------ node_a = NodeId("node-a") node_b = NodeId("node-b") connection = Connection( source=node_a, sink=node_b, edge=SocketConnection( sink_multiaddr=Multiaddr(address="/ip4/127.0.0.1/tcp/10001"), ), ) state = State() state.topology.add_connection(connection) json_repr = state.model_dump_json() restored_state = State.model_validate_json(json_repr) assert ( state.topology.to_snapshot().nodes == restored_state.topology.to_snapshot().nodes ) assert set(state.topology.to_snapshot().connections) == set( restored_state.topology.to_snapshot().connections ) assert restored_state.model_dump_json() == json_repr ================================================ FILE: src/exo/shared/tests/test_xdg_paths.py ================================================ """Tests for XDG Base Directory Specification compliance.""" import os import sys from pathlib import Path from unittest import mock def test_xdg_paths_on_linux(): """Test that XDG paths are used on Linux when XDG env vars are set.""" with ( mock.patch.dict( os.environ, { "XDG_CONFIG_HOME": "/tmp/test-config", "XDG_DATA_HOME": "/tmp/test-data", "XDG_CACHE_HOME": "/tmp/test-cache", }, clear=False, ), mock.patch.object(sys, "platform", "linux"), ): # Re-import to pick up mocked values import importlib import exo.shared.constants as constants importlib.reload(constants) assert Path("/tmp/test-config/exo") == constants.EXO_CONFIG_HOME assert Path("/tmp/test-data/exo") == constants.EXO_DATA_HOME assert Path("/tmp/test-cache/exo") == constants.EXO_CACHE_HOME def test_xdg_default_paths_on_linux(): """Test that XDG default paths are used on Linux when env vars are not set.""" # Remove XDG env vars and EXO_HOME env = { k: v for k, v in os.environ.items() if not k.startswith("XDG_") and k != "EXO_HOME" } with ( mock.patch.dict(os.environ, env, clear=True), mock.patch.object(sys, "platform", "linux"), ): import importlib import exo.shared.constants as constants importlib.reload(constants) home = Path.home() assert home / ".config" / "exo" == constants.EXO_CONFIG_HOME assert home / ".local/share" / "exo" == constants.EXO_DATA_HOME assert home / ".cache" / "exo" == constants.EXO_CACHE_HOME def test_legacy_exo_home_takes_precedence(): """Test that EXO_HOME environment variable takes precedence for backward compatibility.""" with mock.patch.dict( os.environ, { "EXO_HOME": ".custom-exo", "XDG_CONFIG_HOME": "/tmp/test-config", }, clear=False, ): import importlib import exo.shared.constants as constants importlib.reload(constants) home = Path.home() assert home / ".custom-exo" == constants.EXO_CONFIG_HOME assert home / ".custom-exo" == constants.EXO_DATA_HOME def test_macos_uses_traditional_paths(): """Test that macOS uses traditional ~/.exo directory.""" # Remove EXO_HOME to ensure we test the default behavior env = {k: v for k, v in os.environ.items() if k != "EXO_HOME"} with ( mock.patch.dict(os.environ, env, clear=True), mock.patch.object(sys, "platform", "darwin"), ): import importlib import exo.shared.constants as constants importlib.reload(constants) home = Path.home() assert home / ".exo" == constants.EXO_CONFIG_HOME assert home / ".exo" == constants.EXO_DATA_HOME assert home / ".exo" == constants.EXO_CACHE_HOME def test_node_id_in_config_dir(): """Test that node ID keypair is in the config directory.""" import exo.shared.constants as constants assert constants.EXO_NODE_ID_KEYPAIR.parent == constants.EXO_CONFIG_HOME def test_models_in_data_dir(): """Test that models directory is in the data directory.""" # Clear EXO_MODELS_DIR to test default behavior env = {k: v for k, v in os.environ.items() if k != "EXO_MODELS_DIR"} with mock.patch.dict(os.environ, env, clear=True): import importlib import exo.shared.constants as constants importlib.reload(constants) assert constants.EXO_MODELS_DIR.parent == constants.EXO_DATA_HOME ================================================ FILE: src/exo/shared/topology.py ================================================ import contextlib from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from typing import Iterable import rustworkx as rx from pydantic import BaseModel, ConfigDict from exo.shared.types.common import NodeId from exo.shared.types.profiling import ( InterfaceType, NodeNetworkInfo, ThunderboltBridgeStatus, ) from exo.shared.types.topology import ( Connection, Cycle, RDMAConnection, SocketConnection, ) class TopologySnapshot(BaseModel): nodes: Sequence[NodeId] connections: Mapping[ NodeId, Mapping[NodeId, Sequence[SocketConnection | RDMAConnection]] ] model_config = ConfigDict(frozen=True, extra="forbid") @dataclass class Topology: _graph: rx.PyDiGraph[NodeId, SocketConnection | RDMAConnection] = field( init=False, default_factory=rx.PyDiGraph ) _vertex_indices: dict[NodeId, int] = field(init=False, default_factory=dict) def to_snapshot(self) -> TopologySnapshot: return TopologySnapshot( nodes=list(self.list_nodes()), connections=self.map_connections() ) @classmethod def from_snapshot(cls, snapshot: TopologySnapshot) -> "Topology": topology = cls() for node_id in snapshot.nodes: with contextlib.suppress(ValueError): topology.add_node(node_id) for source in snapshot.connections: for sink in snapshot.connections[source]: for edge in snapshot.connections[source][sink]: topology.add_connection( Connection(source=source, sink=sink, edge=edge) ) return topology def add_node(self, node_id: NodeId) -> None: if node_id in self._vertex_indices: return rx_id = self._graph.add_node(node_id) self._vertex_indices[node_id] = rx_id def node_is_leaf(self, node_id: NodeId) -> bool: return ( node_id in self._vertex_indices and len(self._graph.neighbors(self._vertex_indices[node_id])) <= 1 ) def neighbours(self, node_id: NodeId) -> list[NodeId]: return [ self._graph[rx_id] for rx_id in self._graph.neighbors(self._vertex_indices[node_id]) ] def out_edges(self, node_id: NodeId) -> Iterable[Connection]: if node_id not in self._vertex_indices: return [] return ( Connection(source=self._graph[source], sink=self._graph[sink], edge=edge) for source, sink, edge in self._graph.out_edges( self._vertex_indices[node_id] ) ) def contains_node(self, node_id: NodeId) -> bool: return node_id in self._vertex_indices def add_connection(self, conn: Connection) -> None: source, sink, edge = conn.source, conn.sink, conn.edge del conn if edge in self.get_all_connections_between(source, sink): return if source not in self._vertex_indices: self.add_node(source) if sink not in self._vertex_indices: self.add_node(sink) src_id = self._vertex_indices[source] sink_id = self._vertex_indices[sink] _ = self._graph.add_edge(src_id, sink_id, edge) def get_all_connections_between( self, source: NodeId, sink: NodeId ) -> Iterable[SocketConnection | RDMAConnection]: if source not in self._vertex_indices: return [] if sink not in self._vertex_indices: return [] src_id = self._vertex_indices[source] sink_id = self._vertex_indices[sink] try: return self._graph.get_all_edge_data(src_id, sink_id) except rx.NoEdgeBetweenNodes: return [] def list_nodes(self) -> Iterable[NodeId]: return self._graph.nodes() def map_connections( self, ) -> Mapping[NodeId, Mapping[NodeId, Sequence[SocketConnection | RDMAConnection]]]: base: dict[NodeId, dict[NodeId, list[SocketConnection | RDMAConnection]]] = {} for src_id, sink_id, connection in self._graph.weighted_edge_list(): source = self._graph[src_id] sink = self._graph[sink_id] if source not in base: base[source] = {} if sink not in base[source]: base[source][sink] = [] base[source][sink].append(connection) return base def list_connections( self, ) -> Iterable[Connection]: return ( ( Connection( source=self._graph[src_id], sink=self._graph[sink_id], edge=connection, ) ) for src_id, sink_id, connection in self._graph.weighted_edge_list() ) def remove_node(self, node_id: NodeId) -> None: if node_id not in self._vertex_indices: return rx_idx = self._vertex_indices[node_id] self._graph.remove_node(rx_idx) del self._vertex_indices[node_id] def replace_all_out_rdma_connections( self, source: NodeId, new_connections: Sequence[Connection] ) -> None: for conn_idx in self._graph.out_edge_indices(self._vertex_indices[source]): if isinstance(self._graph.get_edge_data_by_index(conn_idx), RDMAConnection): self._graph.remove_edge_from_index(conn_idx) for conn in new_connections: self.add_connection(conn) def remove_connection(self, conn: Connection) -> None: if ( conn.source not in self._vertex_indices or conn.sink not in self._vertex_indices ): return for conn_idx in self._graph.edge_indices_from_endpoints( self._vertex_indices[conn.source], self._vertex_indices[conn.sink] ): if self._graph.get_edge_data_by_index(conn_idx) == conn.edge: self._graph.remove_edge_from_index(conn_idx) def get_cycles(self) -> list[Cycle]: """Get simple cycles in the graph, including singleton cycles""" cycle_idxs = rx.simple_cycles(self._graph) cycles: list[Cycle] = [] for cycle_idx in cycle_idxs: cycle = Cycle(node_ids=[self._graph[idx] for idx in cycle_idx]) cycles.append(cycle) for node_id in self.list_nodes(): cycles.append(Cycle(node_ids=[node_id])) return cycles def get_rdma_cycles(self) -> list[Cycle]: rdma_edges = [ (u, v, conn) for u, v, conn in self._graph.weighted_edge_list() if isinstance(conn, RDMAConnection) ] rdma_graph: rx.PyDiGraph[NodeId, SocketConnection | RDMAConnection] = ( rx.PyDiGraph() ) rdma_graph.add_nodes_from(self._graph.nodes()) for u, v, conn in rdma_edges: rdma_graph.add_edge(u, v, conn) cycle_idxs = rx.simple_cycles(rdma_graph) cycles: list[Cycle] = [] for cycle_idx in cycle_idxs: cycle = Cycle(node_ids=[rdma_graph[idx] for idx in cycle_idx]) cycles.append(cycle) return cycles def get_subgraph_from_nodes(self, node_ids: list[NodeId]) -> "Topology": topology = Topology() for node_id in node_ids: topology.add_node(node_id) for connection in self.list_connections(): if connection.source in node_ids and connection.sink in node_ids: topology.add_connection(connection) return topology def is_rdma_cycle(self, cycle: Cycle) -> bool: node_idxs = [node for node in cycle] rx_idxs = [self._vertex_indices[idx] for idx in node_idxs] for rid in rx_idxs: for neighbor_rid in self._graph.neighbors(rid): if neighbor_rid not in rx_idxs: continue has_rdma = False for edge in self._graph.get_all_edge_data(rid, neighbor_rid): if isinstance(edge, RDMAConnection): has_rdma = True break if not has_rdma: return False return True def get_thunderbolt_bridge_cycles( self, node_tb_bridge_status: Mapping[NodeId, ThunderboltBridgeStatus], node_network: Mapping[NodeId, NodeNetworkInfo], ) -> list[list[NodeId]]: """ Find cycles in the Thunderbolt topology where all nodes have TB bridge enabled. Only returns cycles with >=2 nodes (2+ machines in a loop), as 1 node doesn't cause the broadcast storm problem. """ enabled_nodes = { node_id for node_id, status in node_tb_bridge_status.items() if status.enabled } if len(enabled_nodes) < 2: return [] thunderbolt_ips = _get_ips_with_interface_type( enabled_nodes, node_network, "thunderbolt" ) # Build subgraph with only TB bridge enabled nodes and thunderbolt connections graph: rx.PyDiGraph[NodeId, SocketConnection | RDMAConnection] = rx.PyDiGraph() node_to_idx: dict[NodeId, int] = {} for node_id in enabled_nodes: if node_id in self._vertex_indices: node_to_idx[node_id] = graph.add_node(node_id) for u, v, conn in self._graph.weighted_edge_list(): source_id, sink_id = self._graph[u], self._graph[v] if source_id not in node_to_idx or sink_id not in node_to_idx: continue # Include connection if it's over a thunderbolt interface if ( isinstance(conn, SocketConnection) and conn.sink_multiaddr.ip_address in thunderbolt_ips ): graph.add_edge(node_to_idx[source_id], node_to_idx[sink_id], conn) if isinstance(conn, RDMAConnection): graph.add_edge(node_to_idx[source_id], node_to_idx[sink_id], conn) return [ [graph[idx] for idx in cycle] for cycle in rx.simple_cycles(graph) if len(cycle) >= 2 ] def _get_ips_with_interface_type( node_ids: set[NodeId], node_network: Mapping[NodeId, NodeNetworkInfo], interface_type: InterfaceType, ) -> set[str]: """Get all IP addresses on interfaces of the specified type for the given nodes.""" ips: set[str] = set() for node_id in node_ids: network_info = node_network.get(node_id, NodeNetworkInfo()) for iface in network_info.interfaces: if iface.interface_type == interface_type: ips.add(iface.ip_address) return ips ================================================ FILE: src/exo/shared/tracing.py ================================================ from __future__ import annotations import json import time from collections import defaultdict from collections.abc import Generator from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass, field from pathlib import Path from typing import cast, final from exo.shared.constants import EXO_TRACING_ENABLED from exo.worker.runner.bootstrap import logger # Context variable to track the current trace category for hierarchical nesting _current_category: ContextVar[str | None] = ContextVar("current_category", default=None) @final @dataclass(frozen=True) class TraceEvent: name: str start_us: int duration_us: int rank: int category: str @final @dataclass class CategoryStats: total_us: int = 0 count: int = 0 min_us: int = 0 max_us: int = 0 def add(self, duration_us: int) -> None: if self.count == 0: self.min_us = duration_us self.max_us = duration_us else: self.min_us = min(self.min_us, duration_us) self.max_us = max(self.max_us, duration_us) self.total_us += duration_us self.count += 1 @property def avg_us(self) -> float: return self.total_us / self.count if self.count > 0 else 0.0 @final @dataclass class TraceStats: total_wall_time_us: int = 0 by_category: dict[str, CategoryStats] = field(default_factory=dict) by_rank: dict[int, dict[str, CategoryStats]] = field(default_factory=dict) # Global trace buffer - each rank accumulates traces here _trace_buffer: list[TraceEvent] = [] def _record_span( name: str, start_us: int, duration_us: int, rank: int, category: str ) -> None: _trace_buffer.append( TraceEvent( name=name, start_us=start_us, duration_us=duration_us, rank=rank, category=category, ) ) @contextmanager def trace( name: str, rank: int, category: str = "compute", ) -> Generator[None, None, None]: """Context manager to trace any operation. Nested traces automatically inherit the parent category, creating hierarchical categories like "sync/compute" or "async/comms". Args: name: Name of the operation (e.g., "recv 0", "send 1", "joint_blocks") rank: This rank's ID category: Category for grouping in trace viewer ("comm", "compute", "step") Example: with trace(f"sync {t}", rank, "sync"): with trace("joint_blocks", rank, "compute"): # Recorded with category "sync/compute" hidden_states = some_computation(...) """ if not EXO_TRACING_ENABLED: yield return # Combine with parent category if nested parent = _current_category.get() full_category = f"{parent}/{category}" if parent else category # Set as current for nested traces token = _current_category.set(full_category) try: start_us = int(time.time() * 1_000_000) start_perf = time.perf_counter() yield duration_us = int((time.perf_counter() - start_perf) * 1_000_000) _record_span(name, start_us, duration_us, rank, full_category) finally: _current_category.reset(token) def get_trace_buffer() -> list[TraceEvent]: return list(_trace_buffer) def clear_trace_buffer() -> None: _trace_buffer.clear() def export_trace(traces: list[TraceEvent], output_path: Path) -> None: trace_events: list[dict[str, object]] = [] for event in traces: # Chrome trace format uses "X" for complete events (with duration) chrome_event: dict[str, object] = { "name": event.name, "cat": event.category, "ph": "X", "ts": event.start_us, "dur": event.duration_us, "pid": 0, "tid": event.rank, "args": {"rank": event.rank}, } trace_events.append(chrome_event) ranks_seen = set(t.rank for t in traces) for rank in ranks_seen: trace_events.append( { "name": "thread_name", "ph": "M", # Metadata event "pid": 0, "tid": rank, "args": {"name": f"Rank {rank}"}, } ) chrome_trace = {"traceEvents": trace_events} try: output_path.parent.mkdir(parents=True, exist_ok=True) with open(output_path, "w") as f: json.dump(chrome_trace, f, indent=2) except OSError as e: logger.warning("Failed to export trace to %s: %s", output_path, e) def load_trace_file(path: Path) -> list[TraceEvent]: with open(path) as f: data = cast(dict[str, list[dict[str, object]]], json.load(f)) events = data.get("traceEvents", []) traces: list[TraceEvent] = [] for event in events: # Skip metadata events if event.get("ph") == "M": continue name = str(event.get("name", "")) category = str(event.get("cat", "")) ts_value = event.get("ts", 0) dur_value = event.get("dur", 0) tid_value = event.get("tid", 0) start_us = int(ts_value) if isinstance(ts_value, (int, float, str)) else 0 duration_us = int(dur_value) if isinstance(dur_value, (int, float, str)) else 0 # Get rank from tid or args rank = int(tid_value) if isinstance(tid_value, (int, float, str)) else 0 args = event.get("args") if isinstance(args, dict): args_dict = cast(dict[str, object], args) rank_from_args = args_dict.get("rank") if isinstance(rank_from_args, (int, float, str)): rank = int(rank_from_args) traces.append( TraceEvent( name=name, start_us=start_us, duration_us=duration_us, rank=rank, category=category, ) ) return traces def compute_stats(traces: list[TraceEvent]) -> TraceStats: stats = TraceStats() if not traces: return stats # Calculate wall time from earliest start to latest end min_start = min(t.start_us for t in traces) max_end = max(t.start_us + t.duration_us for t in traces) stats.total_wall_time_us = max_end - min_start # Initialize nested dicts by_category: dict[str, CategoryStats] = defaultdict(CategoryStats) by_rank: dict[int, dict[str, CategoryStats]] = defaultdict( lambda: defaultdict(CategoryStats) ) for event in traces: # By category by_category[event.category].add(event.duration_us) # By rank and category by_rank[event.rank][event.category].add(event.duration_us) stats.by_category = dict(by_category) stats.by_rank = {k: dict(v) for k, v in by_rank.items()} return stats ================================================ FILE: src/exo/shared/types/__init__.py ================================================ ================================================ FILE: src/exo/shared/types/chunks.py ================================================ from collections.abc import Generator from typing import Any, Literal from exo.api.types import ( FinishReason, GenerationStats, ImageGenerationStats, ToolCallItem, TopLogprobItem, Usage, ) from exo.shared.models.model_cards import ModelId from exo.utils.pydantic_ext import TaggedModel from .common import CommandId class BaseChunk(TaggedModel): model: ModelId class TokenChunk(BaseChunk): text: str token_id: int usage: Usage | None finish_reason: Literal["stop", "length", "content_filter"] | None = None stats: GenerationStats | None = None logprob: float | None = None top_logprobs: list[TopLogprobItem] | None = None is_thinking: bool = False class ErrorChunk(BaseChunk): error_message: str finish_reason: Literal["error"] = "error" class ToolCallChunk(BaseChunk): tool_calls: list[ToolCallItem] usage: Usage | None finish_reason: Literal["tool_calls"] = "tool_calls" stats: GenerationStats | None = None class ImageChunk(BaseChunk): data: str chunk_index: int total_chunks: int image_index: int is_partial: bool = False partial_index: int | None = None total_partials: int | None = None stats: ImageGenerationStats | None = None format: Literal["png", "jpeg", "webp"] | None = None finish_reason: FinishReason | None = None error_message: str | None = None def __repr_args__(self) -> Generator[tuple[str, Any], None, None]: for name, value in super().__repr_args__(): # pyright: ignore[reportAny] if name == "data" and hasattr(value, "__len__"): # pyright: ignore[reportAny] yield name, f"<{len(self.data)} chars>" elif name is not None: yield name, value class InputImageChunk(BaseChunk): command_id: CommandId data: str chunk_index: int total_chunks: int def __repr_args__(self) -> Generator[tuple[str, Any], None, None]: for name, value in super().__repr_args__(): # pyright: ignore[reportAny] if name == "data" and hasattr(value, "__len__"): # pyright: ignore[reportAny] yield name, f"<{len(self.data)} chars>" elif name is not None: yield name, value class PrefillProgressChunk(BaseChunk): """Data class for prefill progress events during streaming.""" processed_tokens: int total_tokens: int GenerationChunk = ( TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk | PrefillProgressChunk ) ================================================ FILE: src/exo/shared/types/commands.py ================================================ from pydantic import Field from exo.api.types import ( ImageEditsTaskParams, ImageGenerationTaskParams, ) from exo.shared.models.model_cards import ModelCard, ModelId from exo.shared.types.chunks import InputImageChunk from exo.shared.types.common import CommandId, NodeId, SystemId from exo.shared.types.text_generation import TextGenerationTaskParams from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta from exo.shared.types.worker.shards import Sharding, ShardMetadata from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel class BaseCommand(TaggedModel): command_id: CommandId = Field(default_factory=CommandId) class TestCommand(BaseCommand): __test__ = False class TextGeneration(BaseCommand): task_params: TextGenerationTaskParams class ImageGeneration(BaseCommand): task_params: ImageGenerationTaskParams class ImageEdits(BaseCommand): task_params: ImageEditsTaskParams class PlaceInstance(BaseCommand): model_card: ModelCard sharding: Sharding instance_meta: InstanceMeta min_nodes: int class CreateInstance(BaseCommand): instance: Instance class DeleteInstance(BaseCommand): instance_id: InstanceId class TaskCancelled(BaseCommand): cancelled_command_id: CommandId class TaskFinished(BaseCommand): finished_command_id: CommandId class SendInputChunk(BaseCommand): """Command to send an input image chunk (converted to event by master).""" chunk: InputImageChunk class RequestEventLog(BaseCommand): since_idx: int class StartDownload(BaseCommand): target_node_id: NodeId shard_metadata: ShardMetadata class DeleteDownload(BaseCommand): target_node_id: NodeId model_id: ModelId class CancelDownload(BaseCommand): target_node_id: NodeId model_id: ModelId DownloadCommand = StartDownload | DeleteDownload | CancelDownload Command = ( TestCommand | RequestEventLog | TextGeneration | ImageGeneration | ImageEdits | PlaceInstance | CreateInstance | DeleteInstance | TaskCancelled | TaskFinished | SendInputChunk ) class ForwarderCommand(CamelCaseModel): origin: SystemId command: Command class ForwarderDownloadCommand(CamelCaseModel): origin: SystemId command: DownloadCommand ================================================ FILE: src/exo/shared/types/common.py ================================================ from typing import Self from uuid import uuid4 from pydantic import GetCoreSchemaHandler, field_validator from pydantic_core import core_schema from exo.utils.pydantic_ext import CamelCaseModel class Id(str): def __new__(cls, value: str | None = None) -> Self: return super().__new__(cls, value or str(uuid4())) @classmethod def __get_pydantic_core_schema__( cls, _source: type, handler: GetCoreSchemaHandler ) -> core_schema.CoreSchema: # Just use a plain string schema return core_schema.no_info_after_validator_function( cls, core_schema.str_schema() ) class NodeId(Id): pass class SystemId(Id): pass class ModelId(Id): def normalize(self) -> str: return self.replace("/", "--") def short(self) -> str: return self.split("/")[-1] class SessionId(CamelCaseModel): master_node_id: NodeId election_clock: int class CommandId(Id): pass class Host(CamelCaseModel): ip: str port: int def __str__(self) -> str: return f"{self.ip}:{self.port}" @field_validator("port") @classmethod def check_port(cls, v: int) -> int: if not (0 <= v <= 65535): raise ValueError("Port must be between 0 and 65535") return v ================================================ FILE: src/exo/shared/types/events.py ================================================ from datetime import datetime from typing import final from pydantic import Field from exo.shared.topology import Connection from exo.shared.types.chunks import GenerationChunk, InputImageChunk from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.worker.downloads import DownloadProgress from exo.shared.types.worker.instances import Instance, InstanceId from exo.shared.types.worker.runners import RunnerId, RunnerStatus from exo.utils.info_gatherer.info_gatherer import GatheredInfo from exo.utils.pydantic_ext import CamelCaseModel, FrozenModel, TaggedModel class EventId(Id): """ Newtype around `ID` """ class BaseEvent(TaggedModel): event_id: EventId = Field(default_factory=EventId) # Internal, for debugging. Please don't rely on this field for anything! _master_time_stamp: None | datetime = None class TestEvent(BaseEvent): __test__ = False class TaskCreated(BaseEvent): task_id: TaskId task: Task class TaskAcknowledged(BaseEvent): task_id: TaskId class TaskDeleted(BaseEvent): task_id: TaskId class TaskStatusUpdated(BaseEvent): task_id: TaskId task_status: TaskStatus class TaskFailed(BaseEvent): task_id: TaskId error_type: str error_message: str class InstanceCreated(BaseEvent): instance: Instance def __eq__(self, other: object) -> bool: if isinstance(other, InstanceCreated): return self.instance == other.instance and self.event_id == other.event_id return False class InstanceDeleted(BaseEvent): instance_id: InstanceId class RunnerStatusUpdated(BaseEvent): runner_id: RunnerId runner_status: RunnerStatus class NodeTimedOut(BaseEvent): node_id: NodeId # TODO: bikeshed this name class NodeGatheredInfo(BaseEvent): node_id: NodeId when: str # this is a manually cast datetime overrode by the master when the event is indexed, rather than the local time on the device info: GatheredInfo class NodeDownloadProgress(BaseEvent): download_progress: DownloadProgress class ChunkGenerated(BaseEvent): command_id: CommandId chunk: GenerationChunk class InputChunkReceived(BaseEvent): command_id: CommandId chunk: InputImageChunk class TopologyEdgeCreated(BaseEvent): conn: Connection class TopologyEdgeDeleted(BaseEvent): conn: Connection @final class TraceEventData(FrozenModel): name: str start_us: int duration_us: int rank: int category: str @final class TracesCollected(BaseEvent): task_id: TaskId rank: int traces: list[TraceEventData] @final class TracesMerged(BaseEvent): task_id: TaskId traces: list[TraceEventData] Event = ( TestEvent | TaskCreated | TaskStatusUpdated | TaskFailed | TaskDeleted | TaskAcknowledged | InstanceCreated | InstanceDeleted | RunnerStatusUpdated | NodeTimedOut | NodeGatheredInfo | NodeDownloadProgress | ChunkGenerated | InputChunkReceived | TopologyEdgeCreated | TopologyEdgeDeleted | TracesCollected | TracesMerged ) class IndexedEvent(CamelCaseModel): """An event indexed by the master, with a globally unique index""" idx: int = Field(ge=0) event: Event class GlobalForwarderEvent(CamelCaseModel): """An event the forwarder will serialize and send over the network""" origin_idx: int = Field(ge=0) origin: NodeId session: SessionId event: Event class LocalForwarderEvent(CamelCaseModel): """An event the forwarder will serialize and send over the network""" origin_idx: int = Field(ge=0) origin: SystemId session: SessionId event: Event ================================================ FILE: src/exo/shared/types/memory.py ================================================ from math import ceil from typing import Self, overload from exo.utils.pydantic_ext import FrozenModel class Memory(FrozenModel): in_bytes: int = 0 @classmethod def from_bytes(cls, val: int) -> Self: """Construct a new Memory object from a number of bytes""" return cls(in_bytes=val) @property def in_kb(self) -> int: """The approximate kilobytes this memory represents, rounded up. Setting this property rounds to the nearest byte.""" return ceil(self.in_bytes / 1024) @in_kb.setter def in_kb(self, val: int): """Set this memorys value in kilobytes.""" self.in_bytes = val * 1024 @classmethod def from_kb(cls, val: int) -> Self: """Construct a new Memory object from a number of kilobytes""" return cls(in_bytes=val * 1024) @classmethod def from_float_kb(cls, val: float) -> Self: """Construct a new Memory object from a number of kilobytes, rounding where appropriate""" return cls(in_bytes=round(val * 1024)) @property def in_mb(self) -> int: """The approximate megabytes this memory represents, rounded to nearest MB. Setting this property rounds to the nearest byte.""" return round(self.in_bytes / (1024**2)) @in_mb.setter def in_mb(self, val: int): """Set the megabytes for this memory.""" self.in_bytes = val * (1024**2) @property def in_float_mb(self) -> float: """The megabytes this memory represents as a float. Setting this property rounds to the nearest byte.""" return self.in_bytes / (1024**2) @in_float_mb.setter def in_float_mb(self, val: float): """Set the megabytes for this memory, rounded to the nearest byte.""" self.in_bytes = round(val * (1024**2)) @classmethod def from_mb(cls, val: float) -> Self: """Construct a new Memory object from a number of megabytes""" return cls(in_bytes=round(val * (1024**2))) @classmethod def from_gb(cls, val: float) -> Self: """Construct a new Memory object from a number of megabytes""" return cls(in_bytes=round(val * (1024**3))) @property def in_gb(self) -> float: """The approximate gigabytes this memory represents.""" return self.in_bytes / (1024**3) def __add__(self, other: object) -> "Memory": if isinstance(other, Memory): return Memory.from_bytes(self.in_bytes + other.in_bytes) return NotImplemented def __radd__(self, other: object) -> "Memory": if other == 0: return self return NotImplemented def __sub__(self, other: object) -> "Memory": if isinstance(other, Memory): return Memory.from_bytes(self.in_bytes - other.in_bytes) return NotImplemented def __mul__(self, other: int | float): return Memory.from_bytes(round(self.in_bytes * other)) def __rmul__(self, other: int | float): return self * other @overload def __truediv__(self, other: "Memory") -> float: ... @overload def __truediv__(self, other: int) -> "Memory": ... @overload def __truediv__(self, other: float) -> "Memory": ... def __truediv__(self, other: object) -> "Memory | float": if isinstance(other, Memory): return self.in_bytes / other.in_bytes if isinstance(other, (int, float)): return Memory.from_bytes(round(self.in_bytes / other)) return NotImplemented def __floordiv__(self, other: object) -> "Memory": if isinstance(other, (int, float)): return Memory.from_bytes(int(self.in_bytes // other)) return NotImplemented def __lt__(self, other: object) -> bool: if isinstance(other, Memory): return self.in_bytes < other.in_bytes return NotImplemented def __le__(self, other: object) -> bool: if isinstance(other, Memory): return self.in_bytes <= other.in_bytes return NotImplemented def __gt__(self, other: object) -> bool: if isinstance(other, Memory): return self.in_bytes > other.in_bytes return NotImplemented def __ge__(self, other: object) -> bool: if isinstance(other, Memory): return self.in_bytes >= other.in_bytes return NotImplemented def __eq__(self, other: object) -> bool: if isinstance(other, Memory): return self.in_bytes == other.in_bytes return NotImplemented def __repr__(self) -> str: return f"Memory.from_bytes({self.in_bytes})" def __str__(self) -> str: if self.in_gb > 2: val = self.in_gb unit = "GiB" elif self.in_mb > 2: val = self.in_mb unit = "MiB" elif self.in_kb > 3: val = self.in_kb unit = "KiB" else: val = self.in_bytes unit = "B" return f"{val:.2f} {unit}".rstrip("0").rstrip(".") + f" {unit}" ================================================ FILE: src/exo/shared/types/mlx.py ================================================ """Shared types for MLX-related functionality.""" from collections.abc import Sequence from mlx import core as mx from mlx import nn as nn from mlx_lm.models.cache import ( ArraysCache, CacheList, KVCache, QuantizedKVCache, RotatingKVCache, ) # This list contains one cache entry per transformer layer KVCacheType = Sequence[ KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList ] # Model is a wrapper function to fix the fact that mlx is not strongly typed in the same way that EXO is. # For example - MLX has no guarantee of the interface that nn.Module will expose. But we need a guarantee that it has a __call__() function class Model(nn.Module): layers: list[nn.Module] def __call__( self, x: mx.array, cache: KVCacheType | None, input_embeddings: mx.array | None = None, ) -> mx.array: ... ================================================ FILE: src/exo/shared/types/multiaddr.py ================================================ import re from typing import ClassVar from pydantic import BaseModel, ConfigDict, computed_field, field_validator class Multiaddr(BaseModel): model_config = ConfigDict(frozen=True) address: str PATTERNS: ClassVar[list[str]] = [ r"^/ip4/(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})(/tcp/(\d{1,5}))?(/p2p/[A-Za-z0-9]+)?$", r"^/ip6/([0-9a-fA-F:]+)(/tcp/(\d{1,5}))?(/p2p/[A-Za-z0-9]+)?$", r"^/dns[46]?/([a-zA-Z0-9.-]+)(/tcp/(\d{1,5}))?(/p2p/[A-Za-z0-9]+)?$", ] @field_validator("address") @classmethod def validate_format(cls, v: str) -> str: if not any(re.match(pattern, v) for pattern in cls.PATTERNS): raise ValueError( f"Invalid multiaddr format: {v}. " "Expected format like /ip4/127.0.0.1/tcp/4001 or /dns/example.com/tcp/443" ) return v @computed_field @property def address_type(self) -> str: for pattern in self.PATTERNS: if re.match(pattern, self.address): return pattern.split("/")[1] raise ValueError(f"Invalid multiaddr format: {self.address}") @property def ipv6_address(self) -> str: match = re.match(r"^/ip6/([0-9a-fA-F:]+)", self.address) if not match: raise ValueError( f"Invalid multiaddr format: {self.address}. Expected format like /ip6/::1/tcp/4001" ) return match.group(1) @property def ipv4_address(self) -> str: match = re.match(r"^/ip4/(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})", self.address) if not match: raise ValueError( f"Invalid multiaddr format: {self.address}. Expected format like /ip4/127.0.0.1/tcp/4001" ) return match.group(1) @computed_field @property def ip_address(self) -> str: return self.ipv4_address if self.address_type == "ip4" else self.ipv6_address @computed_field @property def port(self) -> int: match = re.search(r"/tcp/(\d{1,5})", self.address) if not match: raise ValueError( f"Invalid multiaddr format: {self.address}. Expected format like /ip4/127.0.0.1/tcp/4001" ) return int(match.group(1)) def __str__(self) -> str: return self.address ================================================ FILE: src/exo/shared/types/profiling.py ================================================ import shutil from collections.abc import Sequence from pathlib import Path from typing import Literal, Self import psutil from exo.shared.types.memory import Memory from exo.shared.types.thunderbolt import ThunderboltIdentifier from exo.utils.pydantic_ext import CamelCaseModel class MemoryUsage(CamelCaseModel): ram_total: Memory ram_available: Memory swap_total: Memory swap_available: Memory @classmethod def from_bytes( cls, *, ram_total: int, ram_available: int, swap_total: int, swap_available: int ) -> Self: return cls( ram_total=Memory.from_bytes(ram_total), ram_available=Memory.from_bytes(ram_available), swap_total=Memory.from_bytes(swap_total), swap_available=Memory.from_bytes(swap_available), ) @classmethod def from_psutil(cls, *, override_memory: int | None) -> Self: vm = psutil.virtual_memory() sm = psutil.swap_memory() return cls.from_bytes( ram_total=vm.total, ram_available=vm.available if override_memory is None else override_memory, swap_total=sm.total, swap_available=sm.free, ) class DiskUsage(CamelCaseModel): """Disk space usage for the models directory.""" total: Memory available: Memory @classmethod def from_path(cls, path: Path) -> Self: """Get disk usage stats for the partition containing path.""" total, _used, free = shutil.disk_usage(path) return cls( total=Memory.from_bytes(total), available=Memory.from_bytes(free), ) class SystemPerformanceProfile(CamelCaseModel): # TODO: flops_fp16: float gpu_usage: float = 0.0 temp: float = 0.0 sys_power: float = 0.0 pcpu_usage: float = 0.0 ecpu_usage: float = 0.0 InterfaceType = Literal["wifi", "ethernet", "maybe_ethernet", "thunderbolt", "unknown"] class NetworkInterfaceInfo(CamelCaseModel): name: str ip_address: str interface_type: InterfaceType = "unknown" class NodeIdentity(CamelCaseModel): """Static and slow-changing node identification data.""" model_id: str = "Unknown" chip_id: str = "Unknown" friendly_name: str = "Unknown" os_version: str = "Unknown" os_build_version: str = "Unknown" class NodeNetworkInfo(CamelCaseModel): """Network interface information for a node.""" interfaces: Sequence[NetworkInterfaceInfo] = [] class NodeThunderboltInfo(CamelCaseModel): """Thunderbolt interface identifiers for a node.""" interfaces: Sequence[ThunderboltIdentifier] = [] class NodeRdmaCtlStatus(CamelCaseModel): """Whether RDMA is enabled on this node (via rdma_ctl).""" enabled: bool class ThunderboltBridgeStatus(CamelCaseModel): """Whether the Thunderbolt Bridge network service is enabled on this node.""" enabled: bool exists: bool service_name: str | None = None ================================================ FILE: src/exo/shared/types/state.py ================================================ from collections.abc import Mapping, Sequence from datetime import datetime from typing import Any, cast from pydantic import ConfigDict, Field, field_serializer, field_validator from pydantic.alias_generators import to_camel from exo.shared.topology import Topology, TopologySnapshot from exo.shared.types.common import NodeId from exo.shared.types.profiling import ( DiskUsage, MemoryUsage, NodeIdentity, NodeNetworkInfo, NodeRdmaCtlStatus, NodeThunderboltInfo, SystemPerformanceProfile, ThunderboltBridgeStatus, ) from exo.shared.types.tasks import Task, TaskId from exo.shared.types.worker.downloads import DownloadProgress from exo.shared.types.worker.instances import Instance, InstanceId from exo.shared.types.worker.runners import RunnerId, RunnerStatus from exo.utils.pydantic_ext import CamelCaseModel class State(CamelCaseModel): """Global system state. The :class:`Topology` instance is encoded/decoded via an immutable :class:`~shared.topology.TopologySnapshot` to ensure compatibility with standard JSON serialisation. """ model_config = ConfigDict( alias_generator=to_camel, validate_by_name=True, extra="forbid", # I want to reenable this ASAP, but it's causing an issue with TaskStatus strict=True, arbitrary_types_allowed=True, ) instances: Mapping[InstanceId, Instance] = {} runners: Mapping[RunnerId, RunnerStatus] = {} downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {} tasks: Mapping[TaskId, Task] = {} last_seen: Mapping[NodeId, datetime] = {} topology: Topology = Field(default_factory=Topology) last_event_applied_idx: int = Field(default=-1, ge=-1) # Granular node state mappings (update independently at different frequencies) node_identities: Mapping[NodeId, NodeIdentity] = {} node_memory: Mapping[NodeId, MemoryUsage] = {} node_disk: Mapping[NodeId, DiskUsage] = {} node_system: Mapping[NodeId, SystemPerformanceProfile] = {} node_network: Mapping[NodeId, NodeNetworkInfo] = {} node_thunderbolt: Mapping[NodeId, NodeThunderboltInfo] = {} node_thunderbolt_bridge: Mapping[NodeId, ThunderboltBridgeStatus] = {} node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] = {} # Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes) thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = [] @field_serializer("topology", mode="plain") def _encode_topology(self, value: Topology) -> TopologySnapshot: return value.to_snapshot() @field_validator("topology", mode="before") @classmethod def _deserialize_topology(cls, value: object) -> Topology: # noqa: D401 – Pydantic validator signature """Convert an incoming *value* into a :class:`Topology` instance. Accepts either an already constructed :class:`Topology` or a mapping representing :class:`~shared.topology.TopologySnapshot`. """ if isinstance(value, Topology): return value if isinstance(value, Mapping): # likely a snapshot-dict coming from JSON snapshot = TopologySnapshot(**cast(dict[str, Any], value)) # type: ignore[arg-type] return Topology.from_snapshot(snapshot) raise TypeError("Invalid representation for Topology field in State") ================================================ FILE: src/exo/shared/types/tasks.py ================================================ from enum import Enum from pydantic import Field from exo.api.types import ( ImageEditsTaskParams, ImageGenerationTaskParams, ) from exo.shared.types.common import CommandId, Id from exo.shared.types.text_generation import TextGenerationTaskParams from exo.shared.types.worker.instances import BoundInstance, InstanceId from exo.shared.types.worker.runners import RunnerId from exo.shared.types.worker.shards import ShardMetadata from exo.utils.pydantic_ext import TaggedModel class TaskId(Id): pass CANCEL_ALL_TASKS = TaskId("CANCEL_ALL_TASKS") class TaskStatus(str, Enum): Pending = "Pending" Running = "Running" Complete = "Complete" TimedOut = "TimedOut" Failed = "Failed" Cancelled = "Cancelled" class BaseTask(TaggedModel): task_id: TaskId = Field(default_factory=TaskId) task_status: TaskStatus = Field(default=TaskStatus.Pending) instance_id: InstanceId class CreateRunner(BaseTask): # emitted by Worker bound_instance: BoundInstance class DownloadModel(BaseTask): # emitted by Worker shard_metadata: ShardMetadata class LoadModel(BaseTask): # emitted by Worker pass class ConnectToGroup(BaseTask): # emitted by Worker pass class StartWarmup(BaseTask): # emitted by Worker pass class TextGeneration(BaseTask): # emitted by Master command_id: CommandId task_params: TextGenerationTaskParams error_type: str | None = Field(default=None) error_message: str | None = Field(default=None) class CancelTask(BaseTask): cancelled_task_id: TaskId runner_id: RunnerId class ImageGeneration(BaseTask): # emitted by Master command_id: CommandId task_params: ImageGenerationTaskParams error_type: str | None = Field(default=None) error_message: str | None = Field(default=None) class ImageEdits(BaseTask): # emitted by Master command_id: CommandId task_params: ImageEditsTaskParams error_type: str | None = Field(default=None) error_message: str | None = Field(default=None) class Shutdown(BaseTask): # emitted by Worker runner_id: RunnerId Task = ( CreateRunner | DownloadModel | ConnectToGroup | LoadModel | StartWarmup | TextGeneration | CancelTask | ImageGeneration | ImageEdits | Shutdown ) ================================================ FILE: src/exo/shared/types/text_generation.py ================================================ """Canonical internal type for text generation task parameters. All external API formats (Chat Completions, Claude Messages, OpenAI Responses) are converted to TextGenerationTaskParams at the API boundary via adapters. """ from typing import Any, Literal from pydantic import BaseModel from exo.shared.types.common import ModelId MessageRole = Literal["user", "assistant", "system", "developer"] ReasoningEffort = Literal["none", "minimal", "low", "medium", "high", "xhigh"] def resolve_reasoning_params( reasoning_effort: ReasoningEffort | None, enable_thinking: bool | None, ) -> tuple[ReasoningEffort | None, bool | None]: """ enable_thinking=True -> reasoning_effort="medium" enable_thinking=False -> reasoning_effort="none" reasoning_effort="none" -> enable_thinking=False reasoning_effort= -> enable_thinking=True """ resolved_effort: ReasoningEffort | None = reasoning_effort resolved_thinking: bool | None = enable_thinking if reasoning_effort is None and enable_thinking is not None: resolved_effort = "medium" if enable_thinking else "none" if enable_thinking is None and reasoning_effort is not None: resolved_thinking = reasoning_effort != "none" return resolved_effort, resolved_thinking class InputMessage(BaseModel, frozen=True): """Internal message for text generation pipelines.""" role: MessageRole content: str class TextGenerationTaskParams(BaseModel, frozen=True): """Canonical internal task params for text generation. Every API adapter converts its wire type into this before handing off to the master/worker pipeline. """ model: ModelId input: list[InputMessage] instructions: str | None = None max_output_tokens: int | None = None temperature: float | None = None top_p: float | None = None stream: bool = False tools: list[dict[str, Any]] | None = None bench: bool = False top_k: int | None = None stop: str | list[str] | None = None seed: int | None = None chat_template_messages: list[dict[str, Any]] | None = None reasoning_effort: ReasoningEffort | None = None enable_thinking: bool | None = None logprobs: bool = False top_logprobs: int | None = None min_p: float | None = None repetition_penalty: float | None = None repetition_context_size: int | None = None ================================================ FILE: src/exo/shared/types/thunderbolt.py ================================================ import anyio from pydantic import BaseModel, Field from exo.utils.pydantic_ext import CamelCaseModel class ThunderboltConnection(CamelCaseModel): source_uuid: str sink_uuid: str class ThunderboltIdentifier(CamelCaseModel): rdma_interface: str domain_uuid: str link_speed: str = "" ## Intentionally minimal, only collecting data we care about - there's a lot more class _ReceptacleTag(BaseModel, extra="ignore"): receptacle_id_key: str | None = None current_speed_key: str | None = None class _ConnectivityItem(BaseModel, extra="ignore"): domain_uuid_key: str | None = None class ThunderboltConnectivityData(BaseModel, extra="ignore"): domain_uuid_key: str | None = None items: list[_ConnectivityItem] | None = Field(None, alias="_items") receptacle_1_tag: _ReceptacleTag | None = None def ident(self, ifaces: dict[str, str]) -> ThunderboltIdentifier | None: if ( self.domain_uuid_key is None or self.receptacle_1_tag is None or self.receptacle_1_tag.receptacle_id_key is None ): return tag = f"Thunderbolt {self.receptacle_1_tag.receptacle_id_key}" assert tag in ifaces # doesn't need to be an assertion but im confident # if tag not in ifaces: return None iface = f"rdma_{ifaces[tag]}" return ThunderboltIdentifier( rdma_interface=iface, domain_uuid=self.domain_uuid_key, link_speed=self.receptacle_1_tag.current_speed_key or "", ) def conn(self) -> ThunderboltConnection | None: if self.domain_uuid_key is None or self.items is None: return sink_key = next( ( item.domain_uuid_key for item in self.items if item.domain_uuid_key is not None ), None, ) if sink_key is None: return None return ThunderboltConnection( source_uuid=self.domain_uuid_key, sink_uuid=sink_key ) class ThunderboltConnectivity(BaseModel, extra="ignore"): SPThunderboltDataType: list[ThunderboltConnectivityData] = [] @classmethod async def gather(cls) -> list[ThunderboltConnectivityData] | None: proc = await anyio.run_process( ["system_profiler", "SPThunderboltDataType", "-json"], check=False ) if proc.returncode != 0: return None # Saving you from PascalCase while avoiding too much pydantic return ThunderboltConnectivity.model_validate_json( proc.stdout ).SPThunderboltDataType ================================================ FILE: src/exo/shared/types/topology.py ================================================ from collections.abc import Iterator from dataclasses import dataclass from exo.shared.types.common import NodeId from exo.shared.types.multiaddr import Multiaddr from exo.utils.pydantic_ext import FrozenModel @dataclass(frozen=True) class Cycle: node_ids: list[NodeId] def __len__(self) -> int: return self.node_ids.__len__() def __iter__(self) -> Iterator[NodeId]: return self.node_ids.__iter__() class RDMAConnection(FrozenModel): source_rdma_iface: str sink_rdma_iface: str class SocketConnection(FrozenModel): sink_multiaddr: Multiaddr def __hash__(self): return hash(self.sink_multiaddr.ip_address) class Connection(FrozenModel): source: NodeId sink: NodeId edge: RDMAConnection | SocketConnection ================================================ FILE: src/exo/shared/types/worker/downloads.py ================================================ from datetime import timedelta from typing import Literal from pydantic import BaseModel, ConfigDict, Field, PositiveInt from exo.shared.types.common import NodeId from exo.shared.types.memory import Memory from exo.shared.types.worker.shards import ShardMetadata from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel class DownloadProgressData(CamelCaseModel): total: Memory downloaded: Memory downloaded_this_session: Memory completed_files: int total_files: int speed: float eta_ms: int files: dict[str, "DownloadProgressData"] class BaseDownloadProgress(TaggedModel): node_id: NodeId shard_metadata: ShardMetadata model_directory: str = "" class DownloadPending(BaseDownloadProgress): downloaded: Memory = Memory() total: Memory = Memory() class DownloadCompleted(BaseDownloadProgress): total: Memory read_only: bool = False class DownloadFailed(BaseDownloadProgress): error_message: str class DownloadOngoing(BaseDownloadProgress): download_progress: DownloadProgressData DownloadProgress = ( DownloadPending | DownloadCompleted | DownloadFailed | DownloadOngoing ) class ModelSafetensorsIndexMetadata(BaseModel): total_size: PositiveInt class ModelSafetensorsIndex(BaseModel): metadata: ModelSafetensorsIndexMetadata | None weight_map: dict[str, str] class FileListEntry(BaseModel): type: Literal["file", "directory"] path: str size: int | None = None class RepoFileDownloadProgress(BaseModel): repo_id: str repo_revision: str file_path: str downloaded: Memory downloaded_this_session: Memory total: Memory speed: float eta: timedelta status: Literal["not_started", "in_progress", "complete"] start_time: float model_config = ConfigDict(frozen=True) class RepoDownloadProgress(BaseModel): repo_id: str repo_revision: str shard: ShardMetadata completed_files: int total_files: int downloaded: Memory downloaded_this_session: Memory total: Memory overall_speed: float overall_eta: timedelta status: Literal["not_started", "in_progress", "complete"] file_progress: dict[str, RepoFileDownloadProgress] = Field(default_factory=dict) model_config = ConfigDict(frozen=True) ================================================ FILE: src/exo/shared/types/worker/instances.py ================================================ from enum import Enum from pydantic import model_validator from exo.shared.models.model_cards import ModelTask from exo.shared.types.common import Host, Id, NodeId from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel class InstanceId(Id): pass class InstanceMeta(str, Enum): MlxRing = "MlxRing" MlxJaccl = "MlxJaccl" class BaseInstance(TaggedModel): instance_id: InstanceId shard_assignments: ShardAssignments def shard(self, runner_id: RunnerId) -> ShardMetadata | None: return self.shard_assignments.runner_to_shard.get(runner_id, None) class MlxRingInstance(BaseInstance): hosts_by_node: dict[NodeId, list[Host]] ephemeral_port: int class MlxJacclInstance(BaseInstance): jaccl_devices: list[list[str | None]] jaccl_coordinators: dict[NodeId, str] # TODO: Single node instance Instance = MlxRingInstance | MlxJacclInstance class BoundInstance(CamelCaseModel): instance: Instance bound_runner_id: RunnerId bound_node_id: NodeId @property def bound_shard(self) -> ShardMetadata: shard = self.instance.shard(self.bound_runner_id) assert shard is not None return shard @property def is_image_model(self) -> bool: return ( ModelTask.TextToImage in self.bound_shard.model_card.tasks or ModelTask.ImageToImage in self.bound_shard.model_card.tasks ) @model_validator(mode="after") def validate_shard_exists(self) -> "BoundInstance": assert ( self.bound_runner_id in self.instance.shard_assignments.runner_to_shard ), ( "Bound Instance must be constructed with a runner_id that is in the instances assigned shards" ) return self ================================================ FILE: src/exo/shared/types/worker/runner_response.py ================================================ from collections.abc import Generator from typing import Any, Literal from exo.api.types import ( FinishReason, GenerationStats, ImageGenerationStats, ToolCallItem, TopLogprobItem, Usage, ) from exo.utils.pydantic_ext import TaggedModel class BaseRunnerResponse(TaggedModel): pass class TokenizedResponse(BaseRunnerResponse): prompt_tokens: int class GenerationResponse(BaseRunnerResponse): text: str token: int logprob: float | None = None top_logprobs: list[TopLogprobItem] | None = None finish_reason: FinishReason | None = None stats: GenerationStats | None = None usage: Usage | None is_thinking: bool = False class ImageGenerationResponse(BaseRunnerResponse): image_data: bytes format: Literal["png", "jpeg", "webp"] = "png" stats: ImageGenerationStats | None = None image_index: int = 0 def __repr_args__(self) -> Generator[tuple[str, Any], None, None]: for name, value in super().__repr_args__(): # pyright: ignore[reportAny] if name == "image_data": yield name, f"<{len(self.image_data)} bytes>" elif name is not None: yield name, value class PartialImageResponse(BaseRunnerResponse): image_data: bytes format: Literal["png", "jpeg", "webp"] = "png" partial_index: int total_partials: int image_index: int = 0 def __repr_args__(self) -> Generator[tuple[str, Any], None, None]: for name, value in super().__repr_args__(): # pyright: ignore[reportAny] if name == "image_data": yield name, f"<{len(self.image_data)} bytes>" elif name is not None: yield name, value class ToolCallResponse(BaseRunnerResponse): tool_calls: list[ToolCallItem] usage: Usage | None stats: GenerationStats | None = None class FinishedResponse(BaseRunnerResponse): pass class PrefillProgressResponse(BaseRunnerResponse): processed_tokens: int total_tokens: int ================================================ FILE: src/exo/shared/types/worker/runners.py ================================================ from collections.abc import Mapping from pydantic import model_validator from exo.shared.models.model_cards import ModelId from exo.shared.types.common import Id, NodeId from exo.shared.types.worker.shards import ShardMetadata from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel class RunnerId(Id): pass class RunnerError(Exception): pass class BaseRunnerStatus(TaggedModel): def is_running(self): return isinstance(self, RunnerRunning) class RunnerIdle(BaseRunnerStatus): pass class RunnerConnecting(BaseRunnerStatus): pass class RunnerConnected(BaseRunnerStatus): pass class RunnerLoading(BaseRunnerStatus): layers_loaded: int = 0 total_layers: int = 0 class RunnerLoaded(BaseRunnerStatus): pass class RunnerWarmingUp(BaseRunnerStatus): pass class RunnerReady(BaseRunnerStatus): pass class RunnerRunning(BaseRunnerStatus): pass class RunnerShuttingDown(BaseRunnerStatus): pass class RunnerShutdown(BaseRunnerStatus): pass class RunnerFailed(BaseRunnerStatus): error_message: str | None = None RunnerStatus = ( RunnerIdle | RunnerConnecting | RunnerConnected | RunnerLoading | RunnerLoaded | RunnerWarmingUp | RunnerReady | RunnerRunning | RunnerShuttingDown | RunnerShutdown | RunnerFailed ) class ShardAssignments(CamelCaseModel): model_id: ModelId runner_to_shard: Mapping[RunnerId, ShardMetadata] node_to_runner: Mapping[NodeId, RunnerId] @model_validator(mode="after") def validate_runners_exist(self) -> "ShardAssignments": for runner_id in self.node_to_runner.values(): if runner_id not in self.runner_to_shard: raise ValueError( f"Runner {runner_id} in node_to_runner does not exist in runner_to_shard" ) return self ================================================ FILE: src/exo/shared/types/worker/shards.py ================================================ from enum import Enum from typing import TypeAlias, final from pydantic import Field from exo.shared.models.model_cards import ModelCard from exo.utils.pydantic_ext import TaggedModel class Sharding(str, Enum): Tensor = "Tensor" Pipeline = "Pipeline" class BaseShardMetadata(TaggedModel): """ Defines a specific shard of the model that is ready to be run on a device. Replaces previous `Shard` object. """ model_card: ModelCard device_rank: int world_size: int # Error handling; equivalent to monkey-patch, but we can't monkey-patch runner.py # This is kinda annoying because it allocates memory in the ShardMetadata object. Can be rethought after Shanghai. immediate_exception: bool = False should_timeout: float | None = None start_layer: int = Field(ge=0) end_layer: int = Field(ge=0) n_layers: int = Field(ge=0) @property def is_first_layer(self) -> bool: return self.start_layer == 0 @property def is_last_layer(self) -> bool: return self.end_layer == self.n_layers def __hash__(self) -> int: return hash( ( self.model_card.model_id, self.start_layer, self.end_layer, self.n_layers, self.device_rank, self.world_size, ) ) @final class PipelineShardMetadata(BaseShardMetadata): """ Pipeline parallelism shard meta. Layers are represented as a half-open interval [start_layer, end_layer), where start_layer is inclusive and end_layer is exclusive. """ @final class CfgShardMetadata(BaseShardMetadata): """Shard metadata for CFG-parallel image generation models.""" cfg_rank: int # 0 = positive branch, 1 = negative branch cfg_world_size: int = 2 # Pipeline-relative coordinates (computed at placement time) pipeline_rank: int # rank within the pipeline group (0, 1, 2, ...) pipeline_world_size: int # number of nodes per pipeline group @final class TensorShardMetadata(BaseShardMetadata): pass ShardMetadata: TypeAlias = ( PipelineShardMetadata | CfgShardMetadata | TensorShardMetadata ) ================================================ FILE: src/exo/utils/__init__.py ================================================ from typing import Any, Type from .phantom import PhantomData def ensure_type[T](obj: Any, expected_type: Type[T]) -> T: # type: ignore if not isinstance(obj, expected_type): raise TypeError(f"Expected {expected_type}, got {type(obj)}") # type: ignore return obj def todo[T]( msg: str = "This code has not been implemented yet.", _phantom: PhantomData[T] = None, ) -> T: raise NotImplementedError(msg) ================================================ FILE: src/exo/utils/banner.py ================================================ import logging import os import sys import webbrowser from exo.shared.constants import EXO_CONFIG_HOME logger = logging.getLogger(__name__) _FIRST_RUN_MARKER = EXO_CONFIG_HOME / ".dashboard_opened" def _is_first_run() -> bool: return not _FIRST_RUN_MARKER.exists() def _mark_first_run_done() -> None: _FIRST_RUN_MARKER.parent.mkdir(parents=True, exist_ok=True) _FIRST_RUN_MARKER.touch() def print_startup_banner(port: int) -> None: dashboard_url = f"http://localhost:{port}" first_run = _is_first_run() banner = f""" ╔═══════════════════════════════════════════════════════════════════════╗ ║ ║ ║ ███████╗██╗ ██╗ ██████╗ ║ ║ ██╔════╝╚██╗██╔╝██╔═══██╗ ║ ║ █████╗ ╚███╔╝ ██║ ██║ ║ ║ ██╔══╝ ██╔██╗ ██║ ██║ ║ ║ ███████╗██╔╝ ██╗╚██████╔╝ ║ ║ ╚══════╝╚═╝ ╚═╝ ╚═════╝ ║ ║ ║ ║ Distributed AI Inference Cluster ║ ║ ║ ╚═══════════════════════════════════════════════════════════════════════╝ ╔═══════════════════════════════════════════════════════════════════════╗ ║ ║ ║ 🌐 Dashboard & API Ready ║ ║ ║ ║ {dashboard_url}{" " * (69 - len(dashboard_url))}║ ║ ║ ║ Click the URL above to open the dashboard in your browser ║ ║ ║ ╚═══════════════════════════════════════════════════════════════════════╝ """ print(banner, file=sys.stderr) if first_run: # Skip browser open when running inside the native macOS app — # FirstLaunchPopout.swift handles the auto-open with a countdown. if not os.environ.get("EXO_RUNTIME_DIR"): try: webbrowser.open(dashboard_url) logger.info("First run detected — opening dashboard in browser") except Exception: logger.debug("Could not auto-open browser", exc_info=True) _mark_first_run_done() ================================================ FILE: src/exo/utils/channels.py ================================================ import contextlib import multiprocessing as mp from dataclasses import dataclass, field from math import inf from multiprocessing.synchronize import Event from queue import Empty, Full from types import TracebackType from typing import Any, Self from anyio import ( CapacityLimiter, ClosedResourceError, EndOfStream, WouldBlock, to_thread, ) from anyio.streams.memory import ( MemoryObjectReceiveStream as AnyioReceiver, ) from anyio.streams.memory import ( MemoryObjectSendStream as AnyioSender, ) from anyio.streams.memory import ( MemoryObjectStreamState as AnyioState, ) class Sender[T](AnyioSender[T]): def clone(self) -> "Sender[T]": if self._closed: raise ClosedResourceError return Sender(_state=self._state) def clone_receiver(self) -> "Receiver[T]": """Constructs a Receiver using a Senders shared state - similar to calling Receiver.clone() without needing the receiver""" if self._closed: raise ClosedResourceError return Receiver(_state=self._state) class Receiver[T](AnyioReceiver[T]): def clone(self) -> "Receiver[T]": if self._closed: raise ClosedResourceError return Receiver(_state=self._state) def clone_sender(self) -> Sender[T]: """Constructs a Sender using a Receivers shared state - similar to calling Sender.clone() without needing the sender""" if self._closed: raise ClosedResourceError return Sender(_state=self._state) def collect(self) -> list[T]: """Collect all currently available items from this receiver""" out: list[T] = [] while True: try: item = self.receive_nowait() out.append(item) except WouldBlock: break return out async def receive_at_least(self, n: int) -> list[T]: out: list[T] = [] out.append(await self.receive()) out.extend(self.collect()) while len(out) < n: out.append(await self.receive()) out.extend(self.collect()) return out def __enter__(self) -> Self: return self class _MpEndOfStream: pass class MpState[T]: def __init__(self, max_buffer_size: float): if max_buffer_size == inf: max_buffer_size = 0 assert isinstance(max_buffer_size, int), ( "State should only ever be constructed with an integer or math.inf size." ) self.max_buffer_size: float = max_buffer_size self.buffer: mp.Queue[T | _MpEndOfStream] = mp.Queue(max_buffer_size) self.closed: Event = mp.Event() def __getstate__(self): d = self.__dict__.copy() d.pop("__orig_class__", None) return d @dataclass(eq=False) class MpSender[T]: """ An interprocess channel, mimicing the Anyio structure. It should be noted that none of the clone methods are implemented for simplicity, for now. """ _state: MpState[T] = field() def send_nowait(self, item: T) -> None: if self._state.closed.is_set(): raise ClosedResourceError try: self._state.buffer.put(item, block=False) except Full: raise WouldBlock from None except ValueError as e: print("Unreachable code path - let me know!") raise ClosedResourceError from e def send(self, item: T) -> None: if self._state.closed.is_set(): raise ClosedResourceError try: self.send_nowait(item) except WouldBlock: # put anyway, blocking self._state.buffer.put(item, block=True) async def send_async(self, item: T) -> None: await to_thread.run_sync( self.send, item, limiter=CapacityLimiter(1), abandon_on_cancel=True ) def close(self) -> None: if not self._state.closed.is_set(): self._state.closed.set() with contextlib.suppress(Exception): self._state.buffer.put_nowait(_MpEndOfStream()) self._state.buffer.close() # == unique to Mp channels == def join(self) -> None: """Ensure any queued messages are resolved before continuing""" assert self._state.closed.is_set(), ( "Mp channels must be closed before being joined" ) self._state.buffer.join_thread() # == context manager support == def __enter__(self) -> Self: return self def __exit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None, ) -> None: self.close() def __getstate__(self) -> dict[str, Any]: d = self.__dict__.copy() d.pop("__orig_class__", None) return d @dataclass(eq=False) class MpReceiver[T]: """ An interprocess channel, mimicing the Anyio structure. It should be noted that none of the clone methods are implemented for simplicity, for now. """ _state: MpState[T] = field() def receive_nowait(self) -> T: if self._state.closed.is_set(): raise ClosedResourceError try: item = self._state.buffer.get(block=False) if isinstance(item, _MpEndOfStream): self.close() raise EndOfStream return item except Empty: raise WouldBlock from None except ValueError as e: print("Unreachable code path - let me know!") raise ClosedResourceError from e def receive(self) -> T: try: return self.receive_nowait() except WouldBlock: try: item = self._state.buffer.get() except (TypeError, OSError): # Queue pipe can get closed while we are blocked on get(). # The underlying connection._handle becomes None, causing # TypeError in read(handle, remaining). raise ClosedResourceError from None if isinstance(item, _MpEndOfStream): self.close() raise EndOfStream from None return item async def receive_async(self) -> T: return await to_thread.run_sync( self.receive, limiter=CapacityLimiter(1), abandon_on_cancel=True ) def close(self) -> None: if not self._state.closed.is_set(): self._state.closed.set() with contextlib.suppress(Exception): self._state.buffer.put_nowait(_MpEndOfStream()) self._state.buffer.close() # == unique to Mp channels == def join(self) -> None: """Block until all enqueued messages are drained off our side of the buffer""" assert self._state.closed.is_set(), ( "Mp channels must be closed before being joined" ) self._state.buffer.join_thread() # == iterator support == def __iter__(self) -> Self: return self def __next__(self) -> T: try: return self.receive() except EndOfStream: raise StopIteration from None # == async iterator support == def __aiter__(self) -> Self: return self async def __anext__(self) -> T: try: return await self.receive_async() except EndOfStream: raise StopAsyncIteration from None # == context manager support == def __enter__(self) -> Self: return self def __exit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None, ) -> None: self.close() def collect(self) -> list[T]: """Collect all currently available items from this receiver""" out: list[T] = [] while True: try: item = self.receive_nowait() out.append(item) except WouldBlock: break return out def receive_at_least(self, n: int) -> list[T]: out: list[T] = [] out.append(self.receive()) out.extend(self.collect()) while len(out) < n: out.append(self.receive()) out.extend(self.collect()) return out def __getstate__(self): d = self.__dict__.copy() d.pop("__orig_class__", None) return d class channel[T]: # noqa: N801 """Create a pair of asynchronous channels for communicating within the same process""" def __new__(cls, max_buffer_size: float = inf) -> tuple[Sender[T], Receiver[T]]: if max_buffer_size != inf and not isinstance(max_buffer_size, int): raise ValueError("max_buffer_size must be either an integer or math.inf") state = AnyioState[T](max_buffer_size) return Sender(_state=state), Receiver(_state=state) class mp_channel[T]: # noqa: N801 """Create a pair of synchronous channels for interprocess communication""" # max buffer size uses math.inf to represent an unbounded queue, and 0 to represent a yet unimplemented "unbuffered" queue. def __new__(cls, max_buffer_size: float = inf) -> tuple[MpSender[T], MpReceiver[T]]: if ( max_buffer_size == 0 or max_buffer_size != inf and not isinstance(max_buffer_size, int) ): raise ValueError( "max_buffer_size must be either an integer or math.inf. 0-sized buffers are not supported by multiprocessing" ) state = MpState[T](max_buffer_size) return MpSender(_state=state), MpReceiver(_state=state) ================================================ FILE: src/exo/utils/dashboard_path.py ================================================ import sys from pathlib import Path from typing import cast def find_resources() -> Path: resources = _find_resources_in_repo() or _find_resources_in_bundle() if resources is None: raise FileNotFoundError( "Unable to locate resources. Did you clone the repo properly?" ) return resources def _find_resources_in_repo() -> Path | None: current_module = Path(__file__).resolve() for parent in current_module.parents: build = parent / "resources" if build.is_dir(): return build return None def _find_resources_in_bundle() -> Path | None: frozen_root = cast(str | None, getattr(sys, "_MEIPASS", None)) if frozen_root is None: return None candidate = Path(frozen_root) / "resources" if candidate.is_dir(): return candidate return None def find_dashboard() -> Path: dashboard = _find_dashboard_in_repo() or _find_dashboard_in_bundle() if not dashboard: raise FileNotFoundError( "Unable to locate dashboard assets - you probably forgot to run `cd dashboard && npm install && npm run build && cd ..`" ) return dashboard def _find_dashboard_in_repo() -> Path | None: current_module = Path(__file__).resolve() for parent in current_module.parents: build = parent / "dashboard" / "build" if build.is_dir() and (build / "index.html").exists(): return build return None def _find_dashboard_in_bundle() -> Path | None: frozen_root = cast(str | None, getattr(sys, "_MEIPASS", None)) if frozen_root is None: return None candidate = Path(frozen_root) / "dashboard" if candidate.is_dir(): return candidate return None ================================================ FILE: src/exo/utils/disk_event_log.py ================================================ import contextlib import json from collections import OrderedDict from collections.abc import Iterator from datetime import datetime, timezone from io import BufferedRandom, BufferedReader from pathlib import Path import msgspec import zstandard from loguru import logger from pydantic import TypeAdapter from exo.shared.types.events import Event _EVENT_ADAPTER: TypeAdapter[Event] = TypeAdapter(Event) _HEADER_SIZE = 4 # uint32 big-endian _OFFSET_CACHE_SIZE = 128 _MAX_ARCHIVES = 5 def _serialize_event(event: Event) -> bytes: return msgspec.msgpack.encode(event.model_dump(mode="json")) def _deserialize_event(raw: bytes) -> Event: # Decode msgpack into a Python dict, then re-encode as JSON for Pydantic. # Pydantic's validate_json() uses JSON-mode coercion (e.g. string -> enum) # even under strict=True, whereas validate_python() does not. Going through # JSON is the only way to get correct round-trip deserialization without # disabling strict mode or adding casts everywhere. as_json = json.dumps(msgspec.msgpack.decode(raw, type=dict)) return _EVENT_ADAPTER.validate_json(as_json) def _unpack_header(header: bytes) -> int: return int.from_bytes(header, byteorder="big") def _skip_record(f: BufferedReader) -> bool: """Skip one length-prefixed record. Returns False on EOF.""" header = f.read(_HEADER_SIZE) if len(header) < _HEADER_SIZE: return False f.seek(_unpack_header(header), 1) return True def _read_record(f: BufferedReader) -> Event | None: """Read one length-prefixed record. Returns None on EOF.""" header = f.read(_HEADER_SIZE) if len(header) < _HEADER_SIZE: return None length = _unpack_header(header) payload = f.read(length) if len(payload) < length: return None return _deserialize_event(payload) class DiskEventLog: """Append-only event log backed by a file on disk. On-disk format: sequence of length-prefixed msgpack records. Each record is [4-byte big-endian uint32 length][msgpack payload]. Uses a bounded LRU cache of event index → byte offset for efficient random access without storing an offset per event. """ def __init__(self, directory: Path) -> None: self._directory = directory self._directory.mkdir(parents=True, exist_ok=True) self._active_path = directory / "events.bin" self._offset_cache: OrderedDict[int, int] = OrderedDict() self._count: int = 0 # Rotate stale active file from a previous session/crash if self._active_path.exists(): self._rotate(self._active_path, self._directory) self._file: BufferedRandom = open(self._active_path, "w+b") # noqa: SIM115 def _cache_offset(self, idx: int, offset: int) -> None: self._offset_cache[idx] = offset self._offset_cache.move_to_end(idx) if len(self._offset_cache) > _OFFSET_CACHE_SIZE: self._offset_cache.popitem(last=False) def _seek_to(self, f: BufferedReader, target_idx: int) -> None: """Seek f to the byte offset of event target_idx, using cache or scanning forward.""" if target_idx in self._offset_cache: self._offset_cache.move_to_end(target_idx) f.seek(self._offset_cache[target_idx]) return # Find the highest cached index before target_idx scan_from_idx = 0 scan_from_offset = 0 for cached_idx in self._offset_cache: if cached_idx < target_idx: scan_from_idx = cached_idx scan_from_offset = self._offset_cache[cached_idx] # Scan forward, skipping records f.seek(scan_from_offset) for _ in range(scan_from_idx, target_idx): _skip_record(f) self._cache_offset(target_idx, f.tell()) def append(self, event: Event) -> None: packed = _serialize_event(event) self._file.write(len(packed).to_bytes(_HEADER_SIZE, byteorder="big")) self._file.write(packed) self._count += 1 def read_range(self, start: int, end: int) -> Iterator[Event]: """Yield events from index start (inclusive) to end (exclusive).""" end = min(end, self._count) if start < 0 or end < 0 or start >= end: return self._file.flush() with open(self._active_path, "rb") as f: self._seek_to(f, start) for _ in range(end - start): event = _read_record(f) if event is None: break yield event # Cache where we ended up so the next sequential read is a hit if end < self._count: self._cache_offset(end, f.tell()) def read_all(self) -> Iterator[Event]: """Yield all events from the log one at a time.""" if self._count == 0: return self._file.flush() with open(self._active_path, "rb") as f: for _ in range(self._count): event = _read_record(f) if event is None: break yield event def __len__(self) -> int: return self._count def close(self) -> None: """Close the file and rotate active file to compressed archive.""" if self._file.closed: return self._file.close() if self._active_path.exists() and self._count > 0: self._rotate(self._active_path, self._directory) elif self._active_path.exists(): self._active_path.unlink() @staticmethod def _rotate(source: Path, directory: Path) -> None: """Compress source into a timestamped archive. Keeps at most ``_MAX_ARCHIVES`` compressed copies. Oldest beyond the limit are deleted. """ try: stamp = datetime.now(timezone.utc).strftime("%Y-%m-%d_%H-%M-%S_%f") dest = directory / f"events.{stamp}.bin.zst" compressor = zstandard.ZstdCompressor() with open(source, "rb") as f_in, open(dest, "wb") as f_out: compressor.copy_stream(f_in, f_out) source.unlink() logger.info(f"Rotated event log: {source} -> {dest}") # Prune oldest archives beyond the limit archives = sorted(directory.glob("events.*.bin.zst")) for old in archives[:-_MAX_ARCHIVES]: old.unlink() except Exception as e: logger.opt(exception=e).warning(f"Failed to rotate event log {source}") # Clean up the source even if compression fails with contextlib.suppress(OSError): source.unlink() ================================================ FILE: src/exo/utils/event_buffer.py ================================================ from loguru import logger class OrderedBuffer[T]: """ A buffer that resequences events to ensure their ordering is preserved. Currently this buffer doesn't raise any errors if an event is lost This buffer is NOT thread safe, and is designed to only be polled from one source at a time. """ def __init__(self): self.store: dict[int, T] = {} self.next_idx_to_release: int = 0 def ingest(self, idx: int, t: T): """Ingest a sequence into the buffer""" logger.trace(f"Ingested event {t}") if idx < self.next_idx_to_release: return if idx in self.store: assert self.store[idx] == t, ( "Received different messages with identical indices, probable race condition" ) return self.store[idx] = t def drain(self) -> list[T]: """Drain all available events from the buffer""" ret: list[T] = [] while self.next_idx_to_release in self.store: idx = self.next_idx_to_release event = self.store.pop(idx) ret.append(event) self.next_idx_to_release += 1 logger.trace(f"Releasing event {ret}") return ret def drain_indexed(self) -> list[tuple[int, T]]: """Drain all available events from the buffer""" ret: list[tuple[int, T]] = [] while self.next_idx_to_release in self.store: idx = self.next_idx_to_release event = self.store.pop(idx) ret.append((idx, event)) self.next_idx_to_release += 1 logger.trace(f"Releasing event {ret}") return ret class MultiSourceBuffer[SourceId, T]: """ A buffer that resequences events to ensure their ordering is preserved. Tracks events with multiple sources """ def __init__(self): self.stores: dict[SourceId, OrderedBuffer[T]] = {} def ingest(self, idx: int, t: T, source: SourceId): if source not in self.stores: self.stores[source] = OrderedBuffer() buffer = self.stores[source] buffer.ingest(idx, t) def drain(self) -> list[T]: ret: list[T] = [] for store in self.stores.values(): ret.extend(store.drain()) return ret ================================================ FILE: src/exo/utils/fs.py ================================================ import contextlib import os import pathlib import tempfile from typing import LiteralString type StrPath = str | os.PathLike[str] type BytesPath = bytes | os.PathLike[bytes] type StrOrBytesPath = str | bytes | os.PathLike[str] | os.PathLike[bytes] def delete_if_exists(filename: StrOrBytesPath) -> None: with contextlib.suppress(FileNotFoundError): os.remove(filename) def ensure_parent_directory_exists(filename: StrPath) -> None: """ Ensure the directory containing the file exists (create it if necessary). """ pathlib.Path(filename).parent.mkdir(parents=True, exist_ok=True) def ensure_directory_exists(dirname: StrPath) -> None: """ Ensure the directory exists (create it if necessary). """ pathlib.Path(dirname).mkdir(parents=True, exist_ok=True) def make_temp_path(name: LiteralString) -> str: return os.path.join(tempfile.mkdtemp(), name) ================================================ FILE: src/exo/utils/info_gatherer/__init__.py ================================================ ================================================ FILE: src/exo/utils/info_gatherer/info_gatherer.py ================================================ import os import shutil import sys import tomllib from collections.abc import Sequence from dataclasses import dataclass, field from subprocess import CalledProcessError from typing import Self, cast import anyio from anyio import fail_after, open_process, to_thread from anyio.streams.buffered import BufferedByteReceiveStream from anyio.streams.text import TextReceiveStream from loguru import logger from pydantic import ValidationError from exo.shared.constants import EXO_CONFIG_FILE, EXO_MODELS_DIR from exo.shared.types.memory import Memory from exo.shared.types.profiling import ( DiskUsage, MemoryUsage, NetworkInterfaceInfo, ThunderboltBridgeStatus, ) from exo.shared.types.thunderbolt import ( ThunderboltConnection, ThunderboltConnectivity, ThunderboltIdentifier, ) from exo.utils.channels import Sender from exo.utils.pydantic_ext import TaggedModel from exo.utils.task_group import TaskGroup from .macmon import MacmonMetrics from .system_info import ( get_friendly_name, get_model_and_chip, get_network_interfaces, get_os_build_version, get_os_version, ) IS_DARWIN = sys.platform == "darwin" async def _get_thunderbolt_devices() -> set[str] | None: """Get Thunderbolt interface device names (e.g., en2, en3) from hardware ports. Returns None if the networksetup command fails. """ result = await anyio.run_process( ["networksetup", "-listallhardwareports"], check=False, ) if result.returncode != 0: logger.warning( f"networksetup -listallhardwareports failed with code " f"{result.returncode}: {result.stderr.decode()}" ) return None output = result.stdout.decode() thunderbolt_devices: set[str] = set() current_port: str | None = None for line in output.splitlines(): line = line.strip() if line.startswith("Hardware Port:"): current_port = line.split(":", 1)[1].strip() elif line.startswith("Device:") and current_port: device = line.split(":", 1)[1].strip() if "thunderbolt" in current_port.lower(): thunderbolt_devices.add(device) current_port = None return thunderbolt_devices async def _get_bridge_services() -> dict[str, str] | None: """Get mapping of bridge device -> service name from network service order. Returns None if the networksetup command fails. """ result = await anyio.run_process( ["networksetup", "-listnetworkserviceorder"], check=False, ) if result.returncode != 0: logger.warning( f"networksetup -listnetworkserviceorder failed with code " f"{result.returncode}: {result.stderr.decode()}" ) return None # Parse service order to find bridge devices and their service names # Format: "(1) Service Name\n(Hardware Port: ..., Device: bridge0)\n" service_order_output = result.stdout.decode() bridge_services: dict[str, str] = {} # device -> service name current_service: str | None = None for line in service_order_output.splitlines(): line = line.strip() # Match "(N) Service Name" or "(*) Service Name" (disabled) # but NOT "(Hardware Port: ...)" lines if ( line and line.startswith("(") and ")" in line and not line.startswith("(Hardware Port:") ): paren_end = line.index(")") if paren_end + 2 <= len(line): current_service = line[paren_end + 2 :] # Match "(Hardware Port: ..., Device: bridgeX)" elif current_service and "Device: bridge" in line: # Extract device name from "..., Device: bridge0)" device_start = line.find("Device: ") + len("Device: ") device_end = line.find(")", device_start) if device_end > device_start: device = line[device_start:device_end] bridge_services[device] = current_service return bridge_services async def _get_bridge_members(bridge_device: str) -> set[str]: """Get member interfaces of a bridge device via ifconfig.""" result = await anyio.run_process( ["ifconfig", bridge_device], check=False, ) if result.returncode != 0: logger.debug(f"ifconfig {bridge_device} failed with code {result.returncode}") return set() members: set[str] = set() ifconfig_output = result.stdout.decode() for line in ifconfig_output.splitlines(): line = line.strip() if line.startswith("member:"): parts = line.split() if len(parts) > 1: members.add(parts[1]) return members async def _find_thunderbolt_bridge( bridge_services: dict[str, str], thunderbolt_devices: set[str] ) -> str | None: """Find the service name of a bridge containing Thunderbolt interfaces. Returns the service name if found, None otherwise. """ for bridge_device, service_name in bridge_services.items(): members = await _get_bridge_members(bridge_device) if members & thunderbolt_devices: # intersection is non-empty return service_name return None async def _is_service_enabled(service_name: str) -> bool | None: """Check if a network service is enabled. Returns True if enabled, False if disabled, None on error. """ result = await anyio.run_process( ["networksetup", "-getnetworkserviceenabled", service_name], check=False, ) if result.returncode != 0: logger.warning( f"networksetup -getnetworkserviceenabled '{service_name}' " f"failed with code {result.returncode}: {result.stderr.decode()}" ) return None stdout = result.stdout.decode().strip().lower() return stdout == "enabled" class StaticNodeInformation(TaggedModel): """Node information that should NEVER change, to be gathered once at startup""" model: str chip: str os_version: str os_build_version: str @classmethod async def gather(cls) -> Self: model, chip = await get_model_and_chip() return cls( model=model, chip=chip, os_version=get_os_version(), os_build_version=await get_os_build_version(), ) class NodeNetworkInterfaces(TaggedModel): ifaces: Sequence[NetworkInterfaceInfo] class MacThunderboltIdentifiers(TaggedModel): idents: Sequence[ThunderboltIdentifier] class MacThunderboltConnections(TaggedModel): conns: Sequence[ThunderboltConnection] class RdmaCtlStatus(TaggedModel): enabled: bool @classmethod async def gather(cls) -> Self | None: if not IS_DARWIN or shutil.which("rdma_ctl") is None: return None try: with anyio.fail_after(5): proc = await anyio.run_process(["rdma_ctl", "status"], check=False) except (TimeoutError, OSError): return None if proc.returncode != 0: return None output = proc.stdout.decode("utf-8").lower().strip() if "enabled" in output: return cls(enabled=True) if "disabled" in output: return cls(enabled=False) return None class ThunderboltBridgeInfo(TaggedModel): status: ThunderboltBridgeStatus @classmethod async def gather(cls) -> Self | None: """Check if a Thunderbolt Bridge network service is enabled on this node. Detection approach: 1. Find all Thunderbolt interface devices (en2, en3, etc.) from hardware ports 2. Find bridge devices from network service order (not hardware ports, as bridges may not appear there) 3. Check each bridge's members via ifconfig 4. If a bridge contains Thunderbolt interfaces, it's a Thunderbolt Bridge 5. Check if that network service is enabled """ if not IS_DARWIN: return None def _no_bridge_status() -> Self: return cls( status=ThunderboltBridgeStatus( enabled=False, exists=False, service_name=None ) ) try: tb_devices = await _get_thunderbolt_devices() if tb_devices is None: return _no_bridge_status() bridge_services = await _get_bridge_services() if not bridge_services: return _no_bridge_status() tb_service_name = await _find_thunderbolt_bridge( bridge_services, tb_devices ) if not tb_service_name: return _no_bridge_status() enabled = await _is_service_enabled(tb_service_name) if enabled is None: return cls( status=ThunderboltBridgeStatus( enabled=False, exists=True, service_name=tb_service_name ) ) return cls( status=ThunderboltBridgeStatus( enabled=enabled, exists=True, service_name=tb_service_name, ) ) except Exception as e: logger.warning(f"Failed to gather Thunderbolt Bridge info: {e}") return None class NodeConfig(TaggedModel): """Node configuration from EXO_CONFIG_FILE, reloaded from the file only at startup. Other changes should come in through the API and propagate from there""" @classmethod async def gather(cls) -> Self | None: cfg_file = anyio.Path(EXO_CONFIG_FILE) await cfg_file.parent.mkdir(parents=True, exist_ok=True) await cfg_file.touch(exist_ok=True) async with await cfg_file.open("rb") as f: try: contents = (await f.read()).decode("utf-8") data = tomllib.loads(contents) return cls.model_validate(data) except (tomllib.TOMLDecodeError, UnicodeDecodeError, ValidationError): logger.warning("Invalid config file, skipping...") return None class MiscData(TaggedModel): """Node information that may slowly change that doesn't fall into the other categories""" friendly_name: str @classmethod async def gather(cls) -> Self: return cls(friendly_name=await get_friendly_name()) class NodeDiskUsage(TaggedModel): """Disk space information for the models directory.""" disk_usage: DiskUsage @classmethod async def gather(cls) -> Self: return cls( disk_usage=await to_thread.run_sync( lambda: DiskUsage.from_path(EXO_MODELS_DIR) ) ) async def _gather_iface_map() -> dict[str, str] | None: proc = await anyio.run_process( ["networksetup", "-listallhardwareports"], check=False ) if proc.returncode != 0: return None ports: dict[str, str] = {} port = "" for line in proc.stdout.decode("utf-8").split("\n"): if line.startswith("Hardware Port:"): port = line.split(": ")[1] elif line.startswith("Device:"): ports[port] = line.split(": ")[1] port = "" if "" in ports: del ports[""] return ports GatheredInfo = ( MacmonMetrics | MemoryUsage | NodeNetworkInterfaces | MacThunderboltIdentifiers | MacThunderboltConnections | RdmaCtlStatus | ThunderboltBridgeInfo | NodeConfig | MiscData | StaticNodeInformation | NodeDiskUsage ) @dataclass class InfoGatherer: info_sender: Sender[GatheredInfo] interface_watcher_interval: float | None = 10 misc_poll_interval: float | None = 60 system_profiler_interval: float | None = 5 if IS_DARWIN else None memory_poll_rate: float | None = None if IS_DARWIN else 1 macmon_interval: float | None = 1 if IS_DARWIN else None thunderbolt_bridge_poll_interval: float | None = 10 if IS_DARWIN else None static_info_poll_interval: float | None = 60 rdma_ctl_poll_interval: float | None = 10 if IS_DARWIN else None disk_poll_interval: float | None = 30 _tg: TaskGroup = field(init=False, default_factory=TaskGroup) async def run(self): async with self._tg as tg: if IS_DARWIN: if (macmon_path := shutil.which("macmon")) is not None: tg.start_soon(self._monitor_macmon, macmon_path) else: # macmon not installed — fall back to psutil for memory logger.warning( "macmon not found, falling back to psutil for memory monitoring" ) self.memory_poll_rate = 1 tg.start_soon(self._monitor_system_profiler_thunderbolt_data) tg.start_soon(self._monitor_thunderbolt_bridge_status) tg.start_soon(self._monitor_rdma_ctl_status) tg.start_soon(self._watch_system_info) tg.start_soon(self._monitor_memory_usage) tg.start_soon(self._monitor_misc) tg.start_soon(self._monitor_static_info) tg.start_soon(self._monitor_disk_usage) nc = await NodeConfig.gather() if nc is not None: await self.info_sender.send(nc) def shutdown(self): self._tg.cancel_tasks() async def _monitor_static_info(self): if self.static_info_poll_interval is None: return while True: try: with fail_after(30): await self.info_sender.send(await StaticNodeInformation.gather()) except Exception as e: logger.warning(f"Error gathering static node info: {e}") await anyio.sleep(self.static_info_poll_interval) async def _monitor_misc(self): if self.misc_poll_interval is None: return while True: try: with fail_after(10): await self.info_sender.send(await MiscData.gather()) except Exception as e: logger.warning(f"Error gathering misc data: {e}") await anyio.sleep(self.misc_poll_interval) async def _monitor_system_profiler_thunderbolt_data(self): if self.system_profiler_interval is None: return while True: try: with fail_after(30): iface_map = await _gather_iface_map() if iface_map is None: raise ValueError("Failed to gather interface map") data = await ThunderboltConnectivity.gather() assert data is not None idents = [ it for i in data if (it := i.ident(iface_map)) is not None ] await self.info_sender.send( MacThunderboltIdentifiers(idents=idents) ) conns = [it for i in data if (it := i.conn()) is not None] await self.info_sender.send(MacThunderboltConnections(conns=conns)) except Exception as e: logger.warning(f"Error gathering Thunderbolt data: {e}") await anyio.sleep(self.system_profiler_interval) async def _monitor_memory_usage(self): override_memory_env = os.getenv("OVERRIDE_MEMORY_MB") override_memory: int | None = ( Memory.from_mb(int(override_memory_env)).in_bytes if override_memory_env else None ) if self.memory_poll_rate is None: return while True: try: await self.info_sender.send( MemoryUsage.from_psutil(override_memory=override_memory) ) except Exception as e: logger.warning(f"Error gathering memory usage: {e}") await anyio.sleep(self.memory_poll_rate) async def _watch_system_info(self): if self.interface_watcher_interval is None: return while True: try: with fail_after(10): nics = await get_network_interfaces() await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics)) except Exception as e: logger.warning(f"Error gathering network interfaces: {e}") await anyio.sleep(self.interface_watcher_interval) async def _monitor_thunderbolt_bridge_status(self): if self.thunderbolt_bridge_poll_interval is None: return while True: try: with fail_after(30): curr = await ThunderboltBridgeInfo.gather() if curr is not None: await self.info_sender.send(curr) except Exception as e: logger.warning(f"Error gathering Thunderbolt Bridge status: {e}") await anyio.sleep(self.thunderbolt_bridge_poll_interval) async def _monitor_rdma_ctl_status(self): if self.rdma_ctl_poll_interval is None: return while True: try: curr = await RdmaCtlStatus.gather() if curr is not None: await self.info_sender.send(curr) except Exception as e: logger.warning(f"Error gathering RDMA ctl status: {e}") await anyio.sleep(self.rdma_ctl_poll_interval) async def _monitor_disk_usage(self): if self.disk_poll_interval is None: return while True: try: with fail_after(5): await self.info_sender.send(await NodeDiskUsage.gather()) except Exception as e: logger.warning(f"Error gathering disk usage: {e}") await anyio.sleep(self.disk_poll_interval) async def _monitor_macmon(self, macmon_path: str): if self.macmon_interval is None: return # macmon pipe --interval [interval in ms] # Timeout: if macmon produces no output for this many seconds, restart it. # macmon writes every macmon_interval seconds, so 10x that is generous. read_timeout = max(self.macmon_interval * 10, 30) while True: try: async with await open_process( [ macmon_path, "pipe", "--interval", str(self.macmon_interval * 1000), ] ) as p: if not p.stdout: logger.critical("MacMon closed stdout") return stream = TextReceiveStream(BufferedByteReceiveStream(p.stdout)) while True: with fail_after(read_timeout): text = await stream.receive() await self.info_sender.send(MacmonMetrics.from_raw_json(text)) except TimeoutError: logger.warning( f"MacMon produced no output for {read_timeout}s, restarting" ) except CalledProcessError as e: stderr_msg = "no stderr" stderr_output = cast(bytes | str | None, e.stderr) if stderr_output is not None: stderr_msg = ( stderr_output.decode() if isinstance(stderr_output, bytes) else str(stderr_output) ) logger.warning( f"MacMon failed with return code {e.returncode}: {stderr_msg}" ) except Exception as e: logger.warning(f"Error in macmon monitor: {e}") await anyio.sleep(self.macmon_interval) ================================================ FILE: src/exo/utils/info_gatherer/macmon.py ================================================ from typing import Self from pydantic import BaseModel from exo.shared.types.profiling import MemoryUsage, SystemPerformanceProfile from exo.utils.pydantic_ext import TaggedModel class _TempMetrics(BaseModel, extra="ignore"): """Temperature-related metrics returned by macmon.""" cpu_temp_avg: float gpu_temp_avg: float class _MemoryMetrics(BaseModel, extra="ignore"): """Memory-related metrics returned by macmon.""" ram_total: int ram_usage: int swap_total: int swap_usage: int class RawMacmonMetrics(BaseModel, extra="ignore"): """Complete set of metrics returned by macmon. Unknown fields are ignored for forward-compatibility. """ timestamp: str # ignored temp: _TempMetrics memory: _MemoryMetrics ecpu_usage: tuple[int, float] # freq mhz, usage % pcpu_usage: tuple[int, float] # freq mhz, usage % gpu_usage: tuple[int, float] # freq mhz, usage % all_power: float ane_power: float cpu_power: float gpu_power: float gpu_ram_power: float ram_power: float sys_power: float class MacmonMetrics(TaggedModel): system_profile: SystemPerformanceProfile memory: MemoryUsage @classmethod def from_raw(cls, raw: RawMacmonMetrics) -> Self: return cls( system_profile=SystemPerformanceProfile( gpu_usage=raw.gpu_usage[1], temp=raw.temp.gpu_temp_avg, sys_power=raw.sys_power, pcpu_usage=raw.pcpu_usage[1], ecpu_usage=raw.ecpu_usage[1], ), memory=MemoryUsage.from_bytes( ram_total=raw.memory.ram_total, ram_available=(raw.memory.ram_total - raw.memory.ram_usage), swap_total=raw.memory.swap_total, swap_available=(raw.memory.swap_total - raw.memory.swap_usage), ), ) @classmethod def from_raw_json(cls, json: str) -> Self: return cls.from_raw(RawMacmonMetrics.model_validate_json(json)) ================================================ FILE: src/exo/utils/info_gatherer/net_profile.py ================================================ from collections import defaultdict from collections.abc import AsyncGenerator, Mapping import anyio import httpx from anyio import create_task_group from loguru import logger from exo.shared.topology import Topology from exo.shared.types.common import NodeId from exo.shared.types.profiling import NodeNetworkInfo from exo.utils.channels import Sender, channel REACHABILITY_ATTEMPTS = 3 async def check_reachability( target_ip: str, expected_node_id: NodeId, out: dict[NodeId, set[str]], client: httpx.AsyncClient, ) -> None: """Check if a node is reachable at the given IP and verify its identity.""" if ":" in target_ip: # TODO: use real IpAddress types url = f"http://[{target_ip}]:52415/node_id" else: url = f"http://{target_ip}:52415/node_id" remote_node_id = None last_error = None for _ in range(REACHABILITY_ATTEMPTS): try: r = await client.get(url) if r.status_code != 200: await anyio.sleep(1) continue body = r.text.strip().strip('"') if not body: await anyio.sleep(1) continue remote_node_id = NodeId(body) break # expected failure cases except ( httpx.TimeoutException, httpx.NetworkError, ): await anyio.sleep(1) # other failures should be logged on last attempt except httpx.HTTPError as e: last_error = e await anyio.sleep(1) if last_error is not None: logger.warning( f"connect error {type(last_error).__name__} from {target_ip} after {REACHABILITY_ATTEMPTS} attempts; treating as down" ) if remote_node_id is None: return if remote_node_id != expected_node_id: logger.debug( f"Discovered node with unexpected node_id; " f"ip={target_ip}, expected_node_id={expected_node_id}, " f"remote_node_id={remote_node_id}" ) return if remote_node_id not in out: out[remote_node_id] = set() out[remote_node_id].add(target_ip) async def check_reachable( topology: Topology, self_node_id: NodeId, node_network: Mapping[NodeId, NodeNetworkInfo], ) -> AsyncGenerator[tuple[str, NodeId], None]: """Yield (ip, node_id) pairs as reachability probes complete.""" send, recv = channel[tuple[str, NodeId]]() # these are intentionally httpx's defaults so we can tune them later timeout = httpx.Timeout(timeout=5.0) limits = httpx.Limits( max_connections=100, max_keepalive_connections=20, keepalive_expiry=5, ) async def _probe( target_ip: str, expected_node_id: NodeId, client: httpx.AsyncClient, send: Sender[tuple[str, NodeId]], ) -> None: async with send: out: defaultdict[NodeId, set[str]] = defaultdict(set) await check_reachability(target_ip, expected_node_id, out, client) if expected_node_id in out: await send.send((target_ip, expected_node_id)) async with ( httpx.AsyncClient(timeout=timeout, limits=limits, verify=False) as client, create_task_group() as tg, ): for node_id in topology.list_nodes(): if node_id not in node_network: continue if node_id == self_node_id: continue for iface in node_network[node_id].interfaces: tg.start_soon(_probe, iface.ip_address, node_id, client, send.clone()) send.close() with recv: async for item in recv: yield item ================================================ FILE: src/exo/utils/info_gatherer/system_info.py ================================================ import platform import socket import sys from subprocess import CalledProcessError import psutil from anyio import run_process from exo.shared.types.profiling import InterfaceType, NetworkInterfaceInfo def get_os_version() -> str: """Return the OS version string for this node. On macOS this is the macOS version (e.g. ``"15.3"``). On other platforms it falls back to the platform name (e.g. ``"Linux"``). """ if sys.platform == "darwin": version = platform.mac_ver()[0] return version if version else "Unknown" return platform.system() or "Unknown" async def get_os_build_version() -> str: """Return the macOS build version string (e.g. ``"24D5055b"``). On non-macOS platforms, returns ``"Unknown"``. """ if sys.platform != "darwin": return "Unknown" try: process = await run_process(["sw_vers", "-buildVersion"]) except CalledProcessError: return "Unknown" return process.stdout.decode("utf-8", errors="replace").strip() or "Unknown" async def get_friendly_name() -> str: """ Asynchronously gets the 'Computer Name' (friendly name) of a Mac. e.g., "John's MacBook Pro" Returns the name as a string, or None if an error occurs or not on macOS. """ hostname = socket.gethostname() if sys.platform != "darwin": return hostname try: process = await run_process(["scutil", "--get", "ComputerName"]) except CalledProcessError: return hostname return process.stdout.decode("utf-8", errors="replace").strip() or hostname async def _get_interface_types_from_networksetup() -> dict[str, InterfaceType]: """Parse networksetup -listallhardwareports to get interface types.""" if sys.platform != "darwin": return {} try: result = await run_process(["networksetup", "-listallhardwareports"]) except CalledProcessError: return {} types: dict[str, InterfaceType] = {} current_type: InterfaceType = "unknown" for line in result.stdout.decode().splitlines(): if line.startswith("Hardware Port:"): port_name = line.split(":", 1)[1].strip() if "Wi-Fi" in port_name: current_type = "wifi" elif "Ethernet" in port_name or "LAN" in port_name: current_type = "ethernet" elif port_name.startswith("Thunderbolt"): current_type = "thunderbolt" else: current_type = "unknown" elif line.startswith("Device:"): device = line.split(":", 1)[1].strip() # enX is ethernet adapters or thunderbolt - these must be deprioritised if device.startswith("en") and device not in ["en0", "en1"]: current_type = "maybe_ethernet" types[device] = current_type return types async def get_network_interfaces() -> list[NetworkInterfaceInfo]: """ Retrieves detailed network interface information on macOS. Parses output from 'networksetup -listallhardwareports' and 'ifconfig' to determine interface names, IP addresses, and types (ethernet, wifi, vpn, other). Returns a list of NetworkInterfaceInfo objects. """ interfaces_info: list[NetworkInterfaceInfo] = [] interface_types = await _get_interface_types_from_networksetup() for iface, services in psutil.net_if_addrs().items(): for service in services: match service.family: case socket.AF_INET | socket.AF_INET6: interfaces_info.append( NetworkInterfaceInfo( name=iface, ip_address=service.address, interface_type=interface_types.get(iface, "unknown"), ) ) case _: pass return interfaces_info async def get_model_and_chip() -> tuple[str, str]: """Get Mac system information using system_profiler.""" model = "Unknown Model" chip = "Unknown Chip" # TODO: better non mac support if sys.platform != "darwin": return (model, chip) try: process = await run_process( [ "system_profiler", "SPHardwareDataType", ] ) except CalledProcessError: return (model, chip) # less interested in errors here because this value should be hard coded output = process.stdout.decode().strip() model_line = next( (line for line in output.split("\n") if "Model Name" in line), None ) model = model_line.split(": ")[1] if model_line else "Unknown Model" chip_line = next((line for line in output.split("\n") if "Chip" in line), None) chip = chip_line.split(": ")[1] if chip_line else "Unknown Chip" return (model, chip) ================================================ FILE: src/exo/utils/info_gatherer/tests/test_tb_parsing.py ================================================ import sys import pytest from exo.shared.types.thunderbolt import ( ThunderboltConnectivity, ) from exo.utils.info_gatherer.info_gatherer import ( _gather_iface_map, # pyright: ignore[reportPrivateUsage] ) @pytest.mark.anyio @pytest.mark.skipif( sys.platform != "darwin", reason="Thunderbolt info can only be gathered on macos" ) async def test_tb_parsing(): data = await ThunderboltConnectivity.gather() ifaces = await _gather_iface_map() assert ifaces assert data for datum in data: datum.ident(ifaces) datum.conn() ================================================ FILE: src/exo/utils/keyed_backoff.py ================================================ import time from typing import Generic, TypeVar K = TypeVar("K") class KeyedBackoff(Generic[K]): """Tracks exponential backoff state per key.""" def __init__(self, base: float = 0.5, cap: float = 10.0): self._base = base self._cap = cap self._attempts: dict[K, int] = {} self._last_time: dict[K, float] = {} def should_proceed(self, key: K) -> bool: """Returns True if enough time has elapsed since last attempt.""" now = time.monotonic() last = self._last_time.get(key, 0.0) attempts = self._attempts.get(key, 0) delay = min(self._cap, self._base * (2.0**attempts)) return now - last >= delay def record_attempt(self, key: K) -> None: """Record that an attempt was made for this key.""" self._last_time[key] = time.monotonic() self._attempts[key] = self._attempts.get(key, 0) + 1 def reset(self, key: K) -> None: """Reset backoff state for a key (e.g., on success).""" self._attempts.pop(key, None) self._last_time.pop(key, None) ================================================ FILE: src/exo/utils/phantom.py ================================================ class _PhantomData[*T]: """ Internal machinery of the phantom data - it stores nothing. """ type PhantomData[*T] = _PhantomData[*T] | None """ Allows you to use generics in functions without storing anything of that generic type. Just use `None` and you'll be fine """ ================================================ FILE: src/exo/utils/power_sampler.py ================================================ import time from collections import defaultdict from collections.abc import Callable, Mapping from typing import final import anyio from exo.api.types import NodePowerStats, PowerUsage from exo.shared.types.common import NodeId from exo.shared.types.profiling import SystemPerformanceProfile @final class PowerSampler: def __init__( self, get_node_system: Callable[[], Mapping[NodeId, SystemPerformanceProfile]], interval: float = 1.0, ): self._get_node_system = get_node_system self._interval = interval self._samples: defaultdict[NodeId, list[SystemPerformanceProfile]] = ( defaultdict(list) ) self._start_time: float | None = None self._stopped = False def _take_sample(self) -> None: for node_id, profile in self._get_node_system().items(): self._samples[node_id].append(profile) async def run(self) -> None: self._start_time = time.perf_counter() self._take_sample() while not self._stopped: await anyio.sleep(self._interval) self._take_sample() def result(self) -> PowerUsage: self._stopped = True assert self._start_time is not None, "result() called before run()" self._take_sample() elapsed = time.perf_counter() - self._start_time node_stats: list[NodePowerStats] = [] for node_id, profiles in self._samples.items(): n = len(profiles) if n == 0: continue node_stats.append( NodePowerStats( node_id=node_id, samples=n, avg_sys_power=sum(p.sys_power for p in profiles) / n, ) ) total_avg_sys = sum(ns.avg_sys_power for ns in node_stats) return PowerUsage( elapsed_seconds=elapsed, nodes=node_stats, total_avg_sys_power_watts=total_avg_sys, total_energy_joules=total_avg_sys * elapsed, ) ================================================ FILE: src/exo/utils/pydantic_ext.py ================================================ # pyright: reportAny=false, reportUnknownArgumentType=false, reportUnknownVariableType=false from typing import Any, Self from pydantic import BaseModel, ConfigDict, model_serializer, model_validator from pydantic.alias_generators import to_camel from pydantic_core.core_schema import ( SerializerFunctionWrapHandler, ValidatorFunctionWrapHandler, ) class CamelCaseModel(BaseModel): """ A model whose fields are aliased to camel-case from snake-case. """ model_config = ConfigDict( alias_generator=to_camel, validate_by_name=True, extra="forbid", strict=True, ) class FrozenModel(BaseModel): model_config = ConfigDict( alias_generator=to_camel, validate_by_name=True, extra="forbid", strict=True, frozen=True, ) class TaggedModel(CamelCaseModel): @model_serializer(mode="wrap") def _serialize(self, handler: SerializerFunctionWrapHandler): inner = handler(self) return {self.__class__.__name__: inner} @model_validator(mode="wrap") @classmethod def _validate(cls, v: Any, handler: ValidatorFunctionWrapHandler) -> Self: if isinstance(v, dict) and len(v) == 1 and cls.__name__ in v: return handler(v[cls.__name__]) return handler(v) def __str__(self) -> str: return f"{self.__class__.__name__}({super().__str__()})" ================================================ FILE: src/exo/utils/reactive.py ================================================ """ Utilities for reactive variables """ from typing import Protocol class OnChange[T](Protocol): def __call__(self, old_value: T, new_value: T) -> None: ... class Reactive[T]: def __init__(self, initial_value: T, on_change: OnChange[T]): self._value = initial_value self._on_change = on_change @property def value(self): return self._value @value.setter def value(self, new_value: T): old_value = self._value self._value = new_value # don't notify when not changed if old_value == new_value: return # notify of changes self._on_change(old_value=old_value, new_value=new_value) ================================================ FILE: src/exo/utils/task_group.py ================================================ from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from types import TracebackType from typing import Any, Unpack from anyio import create_task_group from anyio.abc import TaskGroup as TaskGroupABC @dataclass class TaskGroup: _tg: TaskGroupABC | None = field(default=None, init=False) _queued: list[tuple[Any, Any, Any]] | None = field(default_factory=list, init=False) def is_running(self) -> bool: return self._tg is not None def cancel_tasks(self): assert self._tg self._tg.cancel_scope.cancel() def cancel_called(self) -> bool: assert self._tg return self._tg.cancel_scope.cancel_called def start_soon[*T]( self, func: Callable[[Unpack[T]], Awaitable[Any]], *args: Unpack[T], name: object = None, ) -> None: assert self._tg is not None assert self._queued is None self._tg.start_soon(func, *args, name=name) def queue[*T]( self, func: Callable[[Unpack[T]], Awaitable[Any]], *args: Unpack[T], name: object = None, ) -> None: assert self._tg is None assert self._queued is not None self._queued.append((func, args, name)) async def __aenter__(self) -> TaskGroupABC: assert self._tg is None assert self._queued is not None self._tg = create_task_group() r = await self._tg.__aenter__() for func, args, name in self._queued: # pyright: ignore[reportAny] self._tg.start_soon(func, *args, name=name) # pyright: ignore[reportAny] self._queued = None return r async def __aexit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None, ) -> bool: """Exit the task group context waiting for all tasks to finish.""" assert self._tg is not None, "aenter sets self.lazy, so it exists when we aexit" assert self._queued is None return await self._tg.__aexit__(exc_type, exc_val, exc_tb) ================================================ FILE: src/exo/utils/tests/test_event_log.py ================================================ from pathlib import Path import pytest from exo.shared.types.events import TestEvent from exo.utils.disk_event_log import DiskEventLog @pytest.fixture def log_dir(tmp_path: Path) -> Path: return tmp_path / "event_log" def test_append_and_read_back(log_dir: Path): log = DiskEventLog(log_dir) events = [TestEvent() for _ in range(5)] for e in events: log.append(e) assert len(log) == 5 result = list(log.read_all()) assert len(result) == 5 for original, restored in zip(events, result, strict=True): assert original.event_id == restored.event_id log.close() def test_read_range(log_dir: Path): log = DiskEventLog(log_dir) events = [TestEvent() for _ in range(10)] for e in events: log.append(e) result = list(log.read_range(3, 7)) assert len(result) == 4 for i, restored in enumerate(result): assert events[3 + i].event_id == restored.event_id log.close() def test_read_range_bounds(log_dir: Path): log = DiskEventLog(log_dir) events = [TestEvent() for _ in range(3)] for e in events: log.append(e) # Start beyond count assert list(log.read_range(5, 10)) == [] # Negative start assert list(log.read_range(-1, 2)) == [] # End beyond count is clamped result = list(log.read_range(1, 100)) assert len(result) == 2 log.close() def test_empty_log(log_dir: Path): log = DiskEventLog(log_dir) assert len(log) == 0 assert list(log.read_all()) == [] assert list(log.read_range(0, 10)) == [] log.close() def _archives(log_dir: Path) -> list[Path]: return sorted(log_dir.glob("events.*.bin.zst")) def test_rotation_on_close(log_dir: Path): log = DiskEventLog(log_dir) log.append(TestEvent()) log.close() active = log_dir / "events.bin" assert not active.exists() archives = _archives(log_dir) assert len(archives) == 1 assert archives[0].stat().st_size > 0 def test_rotation_on_construction_with_stale_file(log_dir: Path): log_dir.mkdir(parents=True, exist_ok=True) (log_dir / "events.bin").write_bytes(b"stale data") log = DiskEventLog(log_dir) archives = _archives(log_dir) assert len(archives) == 1 assert archives[0].exists() assert len(log) == 0 log.close() def test_empty_log_no_archive(log_dir: Path): """Closing an empty log should not leave an archive.""" log = DiskEventLog(log_dir) log.close() active = log_dir / "events.bin" assert not active.exists() assert _archives(log_dir) == [] def test_close_is_idempotent(log_dir: Path): log = DiskEventLog(log_dir) log.append(TestEvent()) log.close() archive = _archives(log_dir) log.close() # should not raise assert _archives(log_dir) == archive def test_successive_sessions(log_dir: Path): """Simulate two master sessions: both archives should be kept.""" log1 = DiskEventLog(log_dir) log1.append(TestEvent()) log1.close() first_archive = _archives(log_dir)[-1] log2 = DiskEventLog(log_dir) log2.append(TestEvent()) log2.append(TestEvent()) log2.close() # Session 1 archive shifted to slot 2, session 2 in slot 1 second_archive = _archives(log_dir)[-1] should_be_first_archive = _archives(log_dir)[-2] assert first_archive.exists() assert second_archive.exists() assert first_archive != second_archive assert should_be_first_archive == first_archive def test_rotation_keeps_at_most_5_archives(log_dir: Path): """After 7 sessions, only the 5 most recent archives should remain.""" all_archives: list[Path] = [] for _ in range(7): log = DiskEventLog(log_dir) log.append(TestEvent()) log.close() all_archives.append(_archives(log_dir)[-1]) for old in all_archives[:2]: assert not old.exists() for recent in all_archives[2:]: assert recent.exists() ================================================ FILE: src/exo/utils/tests/test_mp_channel.py ================================================ import multiprocessing as mp import time import pytest from anyio import fail_after from loguru import logger from exo.utils.channels import MpReceiver, MpSender, mp_channel def foo(recv: MpReceiver[str]): expected = ["hi", "hi 2", "bye"] with recv as r: for item in r: assert item == expected.pop(0) def bar(send: MpSender[str]): logger.warning("hi") send.send("hi") time.sleep(0.1) logger.warning("hi 2") send.send("hi 2") time.sleep(0.1) logger.warning("bye") send.send("bye") time.sleep(0.1) send.close() @pytest.mark.anyio async def test_channel_ipc(): with fail_after(0.5): s, r = mp_channel[str]() p1 = mp.Process(target=foo, args=(r,)) p2 = mp.Process(target=bar, args=(s,)) p1.start() p2.start() p1.join() p2.join() ================================================ FILE: src/exo/utils/tests/test_power_sampler.py ================================================ from collections.abc import Mapping import anyio import pytest from exo.api.types import PowerUsage from exo.shared.types.common import NodeId from exo.shared.types.profiling import SystemPerformanceProfile from exo.utils.power_sampler import PowerSampler def _make_profile(sys_power: float) -> SystemPerformanceProfile: return SystemPerformanceProfile(sys_power=sys_power) NODE_A = NodeId("node-a") NODE_B = NodeId("node-b") @pytest.fixture def single_node_sampler() -> PowerSampler: state: dict[NodeId, SystemPerformanceProfile] = { NODE_A: _make_profile(10.0), } return PowerSampler(get_node_system=lambda: state) @pytest.fixture def multi_node_state() -> dict[NodeId, SystemPerformanceProfile]: return { NODE_A: _make_profile(10.0), NODE_B: _make_profile(20.0), } async def test_single_sample(single_node_sampler: PowerSampler) -> None: """A sampler that runs briefly should capture at least the initial sample.""" async with anyio.create_task_group() as tg: tg.start_soon(single_node_sampler.run) await anyio.sleep(0.05) tg.cancel_scope.cancel() result = single_node_sampler.result() assert len(result.nodes) == 1 assert result.nodes[0].node_id == NODE_A assert result.nodes[0].avg_sys_power == 10.0 assert result.nodes[0].samples >= 1 assert result.elapsed_seconds > 0 async def test_multi_node_averaging( multi_node_state: dict[NodeId, SystemPerformanceProfile], ) -> None: """Power from multiple nodes should be summed for total cluster power.""" sampler = PowerSampler(get_node_system=lambda: multi_node_state) async with anyio.create_task_group() as tg: tg.start_soon(sampler.run) await anyio.sleep(0.05) tg.cancel_scope.cancel() result = sampler.result() assert len(result.nodes) == 2 assert result.total_avg_sys_power_watts == 30.0 async def test_energy_calculation(single_node_sampler: PowerSampler) -> None: """Energy (joules) should be avg_power * elapsed_seconds.""" async with anyio.create_task_group() as tg: tg.start_soon(single_node_sampler.run) await anyio.sleep(0.3) tg.cancel_scope.cancel() result = single_node_sampler.result() expected_energy = result.total_avg_sys_power_watts * result.elapsed_seconds assert result.total_energy_joules == expected_energy async def test_changing_power_is_averaged() -> None: """When power changes mid-sampling, the result should be the average.""" state: dict[NodeId, SystemPerformanceProfile] = { NODE_A: _make_profile(10.0), } sampler = PowerSampler(get_node_system=lambda: state, interval=0.05) async with anyio.create_task_group() as tg: tg.start_soon(sampler.run) await anyio.sleep(0.15) state[NODE_A] = _make_profile(20.0) await anyio.sleep(0.15) tg.cancel_scope.cancel() result = sampler.result() avg = result.nodes[0].avg_sys_power # Should be between 10 and 20, not exactly either assert 10.0 < avg < 20.0 async def test_empty_state() -> None: """A sampler with no nodes should return an empty result.""" empty: Mapping[NodeId, SystemPerformanceProfile] = {} sampler = PowerSampler(get_node_system=lambda: empty) async with anyio.create_task_group() as tg: tg.start_soon(sampler.run) await anyio.sleep(0.05) tg.cancel_scope.cancel() result = sampler.result() assert len(result.nodes) == 0 assert result.total_avg_sys_power_watts == 0.0 assert result.total_energy_joules == 0.0 async def test_result_stops_sampling() -> None: """Calling result() should stop the sampler's run loop.""" state: dict[NodeId, SystemPerformanceProfile] = { NODE_A: _make_profile(10.0), } sampler = PowerSampler(get_node_system=lambda: state, interval=0.02) result: PowerUsage | None = None async with anyio.create_task_group() as tg: tg.start_soon(sampler.run) await anyio.sleep(0.1) result = sampler.result() # run() should exit on its own since _stopped is True await anyio.sleep(0.1) tg.cancel_scope.cancel() assert result is not None assert result.nodes[0].samples >= 2 ================================================ FILE: src/exo/utils/tests/test_tagged.py ================================================ import anyio import pytest from pydantic import BaseModel, TypeAdapter, ValidationError from exo.utils.pydantic_ext import TaggedModel def test_plain_union_prefers_first_member_when_shapes_are_identical(): class Foo1(BaseModel): x: int class Foo2(BaseModel): x: int # Base Pydantic behavior: ambiguous dict goes to the first union member ta = TypeAdapter[Foo1 | Foo2](Foo1 | Foo2) out = ta.validate_python({"x": 1}) assert isinstance(out, Foo1), ( "Base Pydantic should pick the first union member for identical shapes" ) def test_tagged_union_serializes_and_deserializes_two_identical_shapes_correctly(): class Foo1(TaggedModel): x: int class Foo2(TaggedModel): x: int t1 = Foo1(x=1) assert t1.model_dump() == {"Foo1": {"x": 1}} t2 = Foo2(x=2) assert t2.model_dump() == {"Foo2": {"x": 2}} # ---- deserialize (TypeAdapter -> model_validator(before)) ---- ta = TypeAdapter[Foo1 | Foo2](Foo1 | Foo2) out1 = ta.validate_python({"Foo1": {"x": 10}}) assert isinstance(out1, Foo1) and out1.x == 10 out2 = ta.validate_python({"Foo2": {"x": 20}}) assert isinstance(out2, Foo2) and out2.x == 20 def test_tagged_union_rejects_unknown_tag(): class Foo1(TaggedModel): x: int class Foo2(TaggedModel): x: int ta = TypeAdapter[Foo1 | Foo2](Foo1 | Foo2) with pytest.raises(ValidationError): ta.validate_python({"NotARealTag": {"x": 0}}) def test_two_tagged_classes_with_different_shapes_are_independent_and_not_cross_deserializable(): class A1(TaggedModel): x: int class A2(TaggedModel): name: str class B1(TaggedModel): name: str class B2(TaggedModel): active: bool a_payload = A1(x=123).model_dump() b_payload = B1(name="neo").model_dump() assert a_payload == {"A1": {"x": 123}} assert b_payload == {"B1": {"name": "neo"}} ta_a = TypeAdapter[A1 | A2](A1 | A2) ta_b = TypeAdapter[B1 | B2](B1 | B2) with pytest.raises(ValidationError): ta_a.validate_python(b_payload) with pytest.raises(ValidationError): ta_b.validate_python(a_payload) class Inner(TaggedModel): x: int class Outer(TaggedModel): inner: Inner class Wrapper(TaggedModel): outer: Outer label: str class Container(TaggedModel): items: list[Inner] nested: Wrapper def test_single_level_tagging(): inner = Inner(x=10) dumped = inner.model_dump() assert dumped == {"Inner": {"x": 10}} restored = Inner.model_validate(dumped) assert isinstance(restored, Inner) assert restored.x == 10 def test_nested_externally_tagged_union_serializes_recursively(): outer = Outer(inner=Inner(x=42)) dumped = outer.model_dump() assert dumped == {"Outer": {"inner": {"Inner": {"x": 42}}}} restored = Outer.model_validate(dumped) assert isinstance(restored.inner, Inner) assert restored.inner.x == 42 def test_two_level_nested_tagging(): outer = Outer(inner=Inner(x=123)) dumped = outer.model_dump() assert dumped == {"Outer": {"inner": {"Inner": {"x": 123}}}} restored = Outer.model_validate(dumped) assert isinstance(restored.inner, Inner) assert restored.inner.x == 123 def test_three_level_nested_tagging(): wrapper = Wrapper(label="deep", outer=Outer(inner=Inner(x=7))) dumped = wrapper.model_dump() # 3-level structure, each with exactly one tag assert dumped == { "Wrapper": { "label": "deep", "outer": {"Outer": {"inner": {"Inner": {"x": 7}}}}, } } restored = Wrapper.model_validate(dumped) assert isinstance(restored.outer.inner, Inner) assert restored.outer.inner.x == 7 assert restored.label == "deep" def test_lists_and_mixed_nested_structures(): container = Container( items=[Inner(x=1), Inner(x=2)], nested=Wrapper(label="mix", outer=Outer(inner=Inner(x=9))), ) dumped = container.model_dump() assert dumped == { "Container": { "items": [ {"Inner": {"x": 1}}, {"Inner": {"x": 2}}, ], "nested": { "Wrapper": { "label": "mix", "outer": {"Outer": {"inner": {"Inner": {"x": 9}}}}, } }, } } restored = Container.model_validate(dumped) assert isinstance(restored.nested.outer.inner, Inner) assert [i.x for i in restored.items] == [1, 2] def test_no_double_tagging_on_repeated_calls(): """Ensure multiple model_dump calls don't stack tags.""" inner = Inner(x=11) dumped1 = inner.model_dump() dumped2 = inner.model_dump() assert dumped1 == dumped2 == {"Inner": {"x": 11}} outer = Outer(inner=inner) d1 = outer.model_dump() d2 = outer.model_dump() assert d1 == d2 == {"Outer": {"inner": {"Inner": {"x": 11}}}} class L3A(TaggedModel): x: int class L3B(TaggedModel): x: int class L3C(TaggedModel): x: int L3 = L3A | L3B | L3C class L2A(TaggedModel): child: L3 class L2B(TaggedModel): child: L3 class L2C(TaggedModel): child: L3 L2 = L2A | L2B | L2C class L1A(TaggedModel): child: L2 class L1B(TaggedModel): child: L2 class L1C(TaggedModel): child: L2 L1 = L1A | L1B | L1C @pytest.mark.anyio async def test_tagged_union_is_fast(): # payload along the "C" path (worst case for DFS if branches are tried A->B->C) payload = {"L1C": {"child": {"L2C": {"child": {"L3C": {"x": 123}}}}}} with anyio.fail_after(0.1): out = TypeAdapter(L1).validate_python(payload) # type: ignore # Sanity check the result assert out.__class__.__name__ == "L1C" # type: ignore assert out.child.__class__.__name__ == "L2C" # type: ignore assert out.child.child.__class__.__name__ == "L3C" # type: ignore assert out.child.child.x == 123 # type: ignore ================================================ FILE: src/exo/worker/__init__.py ================================================ ================================================ FILE: src/exo/worker/engines/__init__.py ================================================ ================================================ FILE: src/exo/worker/engines/image/__init__.py ================================================ from exo.worker.engines.image.distributed_model import ( DistributedImageModel, initialize_image_model, ) from exo.worker.engines.image.generate import generate_image, warmup_image_generator __all__ = [ "DistributedImageModel", "generate_image", "initialize_image_model", "warmup_image_generator", ] ================================================ FILE: src/exo/worker/engines/image/config.py ================================================ from enum import Enum from pydantic import BaseModel class BlockType(Enum): JOINT = "joint" # Separate image/text streams SINGLE = "single" # Concatenated streams class TransformerBlockConfig(BaseModel): model_config = {"frozen": True} block_type: BlockType count: int has_separate_text_output: bool # True for joint blocks that output text separately class ImageModelConfig(BaseModel): model_family: str block_configs: tuple[TransformerBlockConfig, ...] default_steps: dict[str, int] # {"low": X, "medium": Y, "high": Z} num_sync_steps: int # Number of sync steps for distributed inference guidance_scale: float | None = None # None or <= 1.0 disables CFG @property def total_blocks(self) -> int: return sum(bc.count for bc in self.block_configs) @property def joint_block_count(self) -> int: return sum( bc.count for bc in self.block_configs if bc.block_type == BlockType.JOINT ) @property def single_block_count(self) -> int: return sum( bc.count for bc in self.block_configs if bc.block_type == BlockType.SINGLE ) def get_steps_for_quality(self, quality: str) -> int: return self.default_steps[quality] ================================================ FILE: src/exo/worker/engines/image/distributed_model.py ================================================ from collections.abc import Generator from pathlib import Path from typing import Any, Literal, Optional import mlx.core as mx from mflux.models.common.config.config import Config from PIL import Image from exo.api.types import AdvancedImageParams from exo.download.download_utils import build_model_path from exo.shared.types.worker.instances import BoundInstance from exo.shared.types.worker.shards import CfgShardMetadata, PipelineShardMetadata from exo.worker.engines.image.config import ImageModelConfig from exo.worker.engines.image.models import ( create_adapter_for_model, get_config_for_model, ) from exo.worker.engines.image.models.base import ModelAdapter from exo.worker.engines.image.pipeline import DiffusionRunner from exo.worker.engines.mlx.utils_mlx import mlx_distributed_init, mx_barrier from exo.worker.runner.bootstrap import logger class DistributedImageModel: _config: ImageModelConfig _adapter: ModelAdapter[Any, Any] _runner: DiffusionRunner def __init__( self, model_id: str, local_path: Path, shard_metadata: PipelineShardMetadata | CfgShardMetadata, group: Optional[mx.distributed.Group] = None, quantize: int | None = None, ): config = get_config_for_model(model_id) adapter = create_adapter_for_model(config, model_id, local_path, quantize) has_layer_sharding = ( shard_metadata.start_layer != 0 or shard_metadata.end_layer != shard_metadata.n_layers ) if group is not None and has_layer_sharding: adapter.slice_transformer_blocks( start_layer=shard_metadata.start_layer, end_layer=shard_metadata.end_layer, ) runner = DiffusionRunner( config=config, adapter=adapter, group=group, shard_metadata=shard_metadata, ) if group is not None: logger.info("Initialized distributed diffusion runner") mx.eval(adapter.model.parameters()) # pyright: ignore[reportAny] # TODO(ciaran): Do we need this? mx.eval(adapter.model) # pyright: ignore[reportAny] mx_barrier(group) logger.info(f"Transformer sharded for rank {group.rank()}") else: logger.info("Single-node initialization") self._config = config self._adapter = adapter self._runner = runner @classmethod def from_bound_instance( cls, bound_instance: BoundInstance ) -> "DistributedImageModel": model_id = bound_instance.bound_shard.model_card.model_id model_path = build_model_path(model_id) shard_metadata = bound_instance.bound_shard if not isinstance(shard_metadata, (PipelineShardMetadata, CfgShardMetadata)): raise ValueError( "Expected PipelineShardMetadata or CfgShardMetadata for image generation" ) is_distributed = ( len(bound_instance.instance.shard_assignments.node_to_runner) > 1 ) if is_distributed: logger.info("Starting distributed init for image model") group = mlx_distributed_init(bound_instance) else: group = None return cls( model_id=model_id, local_path=model_path, shard_metadata=shard_metadata, group=group, ) def get_steps_for_quality(self, quality: Literal["low", "medium", "high"]) -> int: """Get the number of inference steps for a quality level.""" return self._config.get_steps_for_quality(quality) def generate( self, prompt: str, height: int, width: int, quality: Literal["low", "medium", "high"] = "medium", seed: int = 2, image_path: Path | None = None, partial_images: int = 0, advanced_params: AdvancedImageParams | None = None, ) -> Generator[Image.Image | tuple[Image.Image, int, int], None, None]: if ( advanced_params is not None and advanced_params.num_inference_steps is not None ): steps = advanced_params.num_inference_steps else: steps = self._config.get_steps_for_quality(quality) guidance_override: float | None = None if advanced_params is not None and advanced_params.guidance is not None: guidance_override = advanced_params.guidance negative_prompt: str | None = None if advanced_params is not None and advanced_params.negative_prompt is not None: negative_prompt = advanced_params.negative_prompt # For edit mode: compute dimensions from input image # This also stores image_paths in the adapter for encode_prompt() if image_path is not None: computed_dims = self._adapter.set_image_dimensions(image_path) if computed_dims is not None: # Override user-provided dimensions with computed ones width, height = computed_dims config = Config( num_inference_steps=steps, height=height, width=width, image_path=image_path, model_config=self._adapter.model.model_config, # pyright: ignore[reportAny] guidance=guidance_override if guidance_override is not None else 4.0, ) if advanced_params is not None and advanced_params.num_sync_steps is not None: num_sync_steps = advanced_params.num_sync_steps else: num_sync_steps = self._config.num_sync_steps for result in self._runner.generate_image( runtime_config=config, prompt=prompt, seed=seed, partial_images=partial_images, guidance_override=guidance_override, negative_prompt=negative_prompt, num_sync_steps=num_sync_steps, ): if isinstance(result, tuple): # Partial image: (GeneratedImage, partial_index, total_partials) image, partial_idx, total_partials = result yield (image, partial_idx, total_partials) else: logger.info("generated image") yield result def initialize_image_model(bound_instance: BoundInstance) -> DistributedImageModel: return DistributedImageModel.from_bound_instance(bound_instance) ================================================ FILE: src/exo/worker/engines/image/generate.py ================================================ import base64 import io import random import tempfile import time from pathlib import Path from typing import Generator, Literal import mlx.core as mx from PIL import Image from exo.api.types import ( AdvancedImageParams, ImageEditsTaskParams, ImageGenerationStats, ImageGenerationTaskParams, ImageSize, ) from exo.shared.types.memory import Memory from exo.shared.types.worker.runner_response import ( ImageGenerationResponse, PartialImageResponse, ) from exo.worker.engines.image.distributed_model import DistributedImageModel def parse_size(size_str: ImageSize) -> tuple[int, int]: """Parse size parameter like '1024x1024' to (width, height) tuple.""" if size_str == "auto": return (1024, 1024) try: parts = size_str.split("x") if len(parts) == 2: width, height = int(parts[0]), int(parts[1]) if width > 0 and height > 0: return (width, height) except (ValueError, AttributeError): pass raise ValueError( f"Invalid size format: '{size_str}'. Expected 'WIDTHxHEIGHT' (e.g., '1024x1024')" ) def warmup_image_generator(model: DistributedImageModel) -> Image.Image | None: """Warmup the image generator with a small image.""" with tempfile.TemporaryDirectory() as tmpdir: # Create a small dummy image for warmup (needed for edit models) dummy_image = Image.new("RGB", (256, 256), color=(128, 128, 128)) dummy_path = Path(tmpdir) / "warmup.png" dummy_image.save(dummy_path) warmup_params = AdvancedImageParams(num_inference_steps=2) for result in model.generate( prompt="Warmup", height=256, width=256, quality="low", image_path=dummy_path, advanced_params=warmup_params, ): if not isinstance(result, tuple): return result return None def generate_image( model: DistributedImageModel, task: ImageGenerationTaskParams | ImageEditsTaskParams, ) -> Generator[ImageGenerationResponse | PartialImageResponse, None, None]: """Generate image(s), optionally yielding partial results. When partial_images > 0 or stream=True, yields PartialImageResponse for intermediate images, then ImageGenerationResponse for the final image. Yields: PartialImageResponse for intermediate images (if partial_images > 0, first image only) ImageGenerationResponse for final complete images """ width, height = parse_size(task.size) quality: Literal["low", "medium", "high"] = task.quality or "medium" advanced_params = task.advanced_params if advanced_params is not None and advanced_params.seed is not None: base_seed = advanced_params.seed else: base_seed = random.randint(0, 2**32 - 1) is_bench = getattr(task, "bench", False) num_images = task.n or 1 generation_start_time: float = 0.0 if is_bench: mx.reset_peak_memory() generation_start_time = time.perf_counter() partial_images = ( task.partial_images if task.partial_images is not None and task.stream is not None and task.stream else 0 ) image_path: Path | None = None with tempfile.TemporaryDirectory() as tmpdir: if isinstance(task, ImageEditsTaskParams): # Decode base64 image data and save to temp file image_path = Path(tmpdir) / "input.png" image_path.write_bytes(base64.b64decode(task.image_data)) if task.size == "auto": with Image.open(image_path) as img: width, height = img.size for image_num in range(num_images): # Increment seed for each image to ensure unique results current_seed = base_seed + image_num for result in model.generate( prompt=task.prompt, height=height, width=width, quality=quality, seed=current_seed, image_path=image_path, partial_images=partial_images, advanced_params=advanced_params, ): if isinstance(result, tuple): # Partial image: (Image, partial_index, total_partials) image, partial_idx, total_partials = result buffer = io.BytesIO() image_format = task.output_format.upper() if image_format == "JPG": image_format = "JPEG" if image_format == "JPEG" and image.mode == "RGBA": image = image.convert("RGB") image.save(buffer, format=image_format) yield PartialImageResponse( image_data=buffer.getvalue(), format=task.output_format, partial_index=partial_idx, total_partials=total_partials, image_index=image_num, ) else: image = result # Only include stats on the final image stats: ImageGenerationStats | None = None if is_bench and image_num == num_images - 1: generation_end_time = time.perf_counter() total_generation_time = ( generation_end_time - generation_start_time ) num_inference_steps = model.get_steps_for_quality(quality) total_steps = num_inference_steps * num_images seconds_per_step = ( total_generation_time / total_steps if total_steps > 0 else 0.0 ) peak_memory = Memory.from_bytes(mx.get_peak_memory()) stats = ImageGenerationStats( seconds_per_step=seconds_per_step, total_generation_time=total_generation_time, num_inference_steps=num_inference_steps, num_images=num_images, image_width=width, image_height=height, peak_memory_usage=peak_memory, ) buffer = io.BytesIO() image_format = task.output_format.upper() if image_format == "JPG": image_format = "JPEG" if image_format == "JPEG" and image.mode == "RGBA": image = image.convert("RGB") image.save(buffer, format=image_format) yield ImageGenerationResponse( image_data=buffer.getvalue(), format=task.output_format, stats=stats, image_index=image_num, ) ================================================ FILE: src/exo/worker/engines/image/models/__init__.py ================================================ from pathlib import Path from typing import Any, Callable from exo.worker.engines.image.config import ImageModelConfig from exo.worker.engines.image.models.base import ModelAdapter from exo.worker.engines.image.models.flux import ( FLUX_DEV_CONFIG, FLUX_KONTEXT_CONFIG, FLUX_SCHNELL_CONFIG, FluxKontextModelAdapter, FluxModelAdapter, ) from exo.worker.engines.image.models.qwen import ( QWEN_IMAGE_CONFIG, QWEN_IMAGE_EDIT_CONFIG, QwenEditModelAdapter, QwenModelAdapter, ) __all__: list[str] = [] # Type alias for adapter factory functions # Factory takes (config, model_id, local_path, quantize) and returns a ModelAdapter AdapterFactory = Callable[ [ImageModelConfig, str, Path, int | None], ModelAdapter[Any, Any] ] # Registry maps model_family string to adapter factory _ADAPTER_REGISTRY: dict[str, AdapterFactory] = { "flux": FluxModelAdapter, "flux-kontext": FluxKontextModelAdapter, "qwen-edit": QwenEditModelAdapter, "qwen": QwenModelAdapter, } # Config registry: maps model ID patterns to configs # Order matters: longer/more-specific patterns must come before shorter ones _CONFIG_REGISTRY: dict[str, ImageModelConfig] = { "flux.1-schnell": FLUX_SCHNELL_CONFIG, "flux.1-kontext": FLUX_KONTEXT_CONFIG, # Must come before "flux.1-dev" for pattern matching "flux.1-krea-dev": FLUX_DEV_CONFIG, # Must come before "flux.1-dev" for pattern matching "flux.1-dev": FLUX_DEV_CONFIG, "qwen-image-edit": QWEN_IMAGE_EDIT_CONFIG, # Must come before "qwen-image" for pattern matching "qwen-image": QWEN_IMAGE_CONFIG, } def get_config_for_model(model_id: str) -> ImageModelConfig: """Get configuration for a model ID. Args: model_id: The model identifier (e.g., "black-forest-labs/FLUX.1-schnell") Returns: The model configuration Raises: ValueError: If no configuration found for model ID """ model_id_lower = model_id.lower() for pattern, config in _CONFIG_REGISTRY.items(): if pattern in model_id_lower: return config raise ValueError(f"No configuration found for model: {model_id}") def create_adapter_for_model( config: ImageModelConfig, model_id: str, local_path: Path, quantize: int | None = None, ) -> ModelAdapter[Any, Any]: """Create a model adapter for the given configuration. Args: config: The model configuration model_id: The model identifier local_path: Path to the model weights quantize: Optional quantization bits Returns: A ModelAdapter instance Raises: ValueError: If no adapter found for model family """ factory = _ADAPTER_REGISTRY.get(config.model_family) if factory is None: raise ValueError(f"No adapter found for model family: {config.model_family}") return factory(config, model_id, local_path, quantize) ================================================ FILE: src/exo/worker/engines/image/models/base.py ================================================ from abc import ABC, abstractmethod from pathlib import Path from typing import TYPE_CHECKING, Any, Generic, TypeVar import mlx.core as mx from mflux.models.common.config.config import Config from mflux.models.common.latent_creator.latent_creator import Img2Img, LatentCreator from mflux.utils.image_util import ImageUtil from PIL import Image from exo.worker.engines.image.config import ImageModelConfig if TYPE_CHECKING: from exo.worker.engines.image.pipeline.block_wrapper import ( JointBlockWrapper, SingleBlockWrapper, ) ModelT = TypeVar("ModelT") TransformerT = TypeVar("TransformerT") RotaryEmbeddings = mx.array | tuple[mx.array, mx.array] class PromptData(ABC): @property @abstractmethod def prompt_embeds(self) -> mx.array: ... @property @abstractmethod def pooled_prompt_embeds(self) -> mx.array: ... @property @abstractmethod def negative_prompt_embeds(self) -> mx.array | None: ... @property @abstractmethod def negative_pooled_prompt_embeds(self) -> mx.array | None: ... @abstractmethod def get_encoder_hidden_states_mask( self, positive: bool = True ) -> mx.array | None: ... @property @abstractmethod def cond_image_grid( self, ) -> tuple[int, int, int] | list[tuple[int, int, int]] | None: """Conditioning image grid dimensions for edit mode. Returns: Grid dimensions (edit) or None (standard generation). """ ... @property @abstractmethod def conditioning_latents(self) -> mx.array | None: """Conditioning latents for edit mode. Returns: Conditioning latents array for image editing, None for standard generation. """ ... @property @abstractmethod def kontext_image_ids(self) -> mx.array | None: """Kontext-style position IDs for image conditioning. For FLUX.1-Kontext models, returns position IDs with first_coord=1 to distinguish conditioning tokens from generation tokens (first_coord=0). Returns: Position IDs array [1, seq_len, 3] for Kontext, None for other models. """ ... @abstractmethod def get_batched_cfg_data( self, ) -> tuple[mx.array, mx.array, mx.array | None, mx.array | None] | None: """Get embeddings for CFG with batch_size=2. Combines positive and negative embeddings into batched tensors for a single forward pass. Pads shorter sequences to max length. Attention mask is used to mask padding. Returns: None if model doesn't support CFG, otherwise tuple of: - batched_embeds: [2, max_seq, hidden] (positive then negative) - batched_mask: [2, max_seq] attention mask - batched_pooled: [2, hidden] pooled embeddings or None - conditioning_latents: [2, latent_seq, latent_dim] or None TODO(ciaran): type this """ ... @abstractmethod def get_cfg_branch_data( self, positive: bool ) -> tuple[mx.array, mx.array | None, mx.array | None, mx.array | None]: """Get embeddings for a single CFG branch (positive or negative). Used for sequential CFG and CFG parallel modes where we process one branch at a time instead of batching. Args: positive: True for positive prompt, False for negative prompt Returns: Tuple of: - embeds: [1, seq, hidden] prompt embeddings - mask: [1, seq] attention mask or None - pooled: [1, hidden] pooled embeddings or None - conditioning_latents: [1, latent_seq, latent_dim] or None """ ... class ModelAdapter(ABC, Generic[ModelT, TransformerT]): _config: ImageModelConfig _model: ModelT _transformer: TransformerT @property def config(self) -> ImageModelConfig: return self._config @property def model(self) -> ModelT: return self._model @property def transformer(self) -> TransformerT: return self._transformer @property @abstractmethod def hidden_dim(self) -> int: ... @property @abstractmethod def needs_cfg(self) -> bool: """Whether this model uses classifier-free guidance.""" ... @abstractmethod def _get_latent_creator(self) -> type: ... @abstractmethod def get_joint_block_wrappers( self, text_seq_len: int, encoder_hidden_states_mask: mx.array | None = None, ) -> list["JointBlockWrapper[Any]"]: """Create wrapped joint transformer blocks with pipefusion support. Args: text_seq_len: Number of text tokens (constant for generation) encoder_hidden_states_mask: Attention mask for text (Qwen only) Returns: List of wrapped joint blocks ready for pipefusion """ ... @abstractmethod def get_single_block_wrappers( self, text_seq_len: int, ) -> list["SingleBlockWrapper[Any]"]: """Create wrapped single transformer blocks with pipefusion support. Args: text_seq_len: Number of text tokens (constant for generation) Returns: List of wrapped single blocks ready for pipefusion """ ... @abstractmethod def slice_transformer_blocks( self, start_layer: int, end_layer: int, ): """Remove transformer blocks outside the assigned range. This should be called BEFORE mx.eval() to avoid loading unused weights in distributed mode. Args: start_layer: First layer index (inclusive) assigned to this node end_layer: Last layer index (exclusive) assigned to this node """ ... def set_image_dimensions(self, image_path: Path) -> tuple[int, int] | None: """Default implementation: no dimension computation needed. Override in edit adapters to compute dimensions from input image. TODO(ciaran): this is a hack Returns: None (use user-specified dimensions) """ return None def create_latents(self, seed: int, runtime_config: Config) -> mx.array: """Create initial latents. Uses model-specific latent creator.""" model: Any = self.model return LatentCreator.create_for_txt2img_or_img2img( seed=seed, height=runtime_config.height, width=runtime_config.width, img2img=Img2Img( vae=model.vae, # pyright: ignore[reportAny] latent_creator=self._get_latent_creator(), sigmas=runtime_config.scheduler.sigmas, # pyright: ignore[reportAny] init_time_step=runtime_config.init_time_step, image_path=runtime_config.image_path, ), ) def decode_latents( self, latents: mx.array, runtime_config: Config, seed: int, prompt: str, ) -> Image.Image: model: Any = self.model # Allow attribute access on model latents = self._get_latent_creator().unpack_latents( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] latents=latents, height=runtime_config.height, width=runtime_config.width, ) decoded = model.vae.decode(latents) # pyright: ignore[reportAny] # TODO(ciaran): # from mflux.models.common.vae.vae_util import VAEUtil # VAEUtil.decode(vae=model.vae, latents=latents, tiling_config=self.tiling_config) generated_image = ImageUtil.to_image( decoded_latents=decoded, # pyright: ignore[reportAny] config=runtime_config, seed=seed, prompt=prompt, quantization=model.bits, # pyright: ignore[reportAny] lora_paths=model.lora_paths, # pyright: ignore[reportAny] lora_scales=model.lora_scales, # pyright: ignore[reportAny] image_path=runtime_config.image_path, image_strength=runtime_config.image_strength, generation_time=0, ) return generated_image.image @abstractmethod def encode_prompt( self, prompt: str, negative_prompt: str | None = None ) -> "PromptData": ... @abstractmethod def compute_embeddings( self, hidden_states: mx.array, prompt_embeds: mx.array, ) -> tuple[mx.array, mx.array]: ... @abstractmethod def compute_text_embeddings( self, t: int, runtime_config: Config, pooled_prompt_embeds: mx.array | None = None, hidden_states: mx.array | None = None, ) -> mx.array: ... @abstractmethod def compute_rotary_embeddings( self, prompt_embeds: mx.array, runtime_config: Config, encoder_hidden_states_mask: mx.array | None = None, cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]] | None = None, kontext_image_ids: mx.array | None = None, ) -> RotaryEmbeddings: ... def merge_streams( self, hidden_states: mx.array, encoder_hidden_states: mx.array, ) -> mx.array: return mx.concatenate([encoder_hidden_states, hidden_states], axis=1) @abstractmethod def apply_guidance( self, noise_positive: mx.array, noise_negative: mx.array, guidance_scale: float, ) -> mx.array: """Apply classifier-free guidance to combine positive/negative predictions. Only called when needs_cfg is True. Args: noise_positive: Noise prediction from positive prompt noise_negative: Noise prediction from negative prompt guidance_scale: Guidance strength Returns: Guided noise prediction """ ... def final_projection( self, hidden_states: mx.array, text_embeddings: mx.array, ) -> mx.array: transformer: Any = self.transformer hidden_states = transformer.norm_out(hidden_states, text_embeddings) # pyright: ignore[reportAny] return transformer.proj_out(hidden_states) # pyright: ignore[reportAny] ================================================ FILE: src/exo/worker/engines/image/models/flux/__init__.py ================================================ from exo.worker.engines.image.models.flux.adapter import FluxModelAdapter from exo.worker.engines.image.models.flux.config import ( FLUX_DEV_CONFIG, FLUX_KONTEXT_CONFIG, FLUX_SCHNELL_CONFIG, ) from exo.worker.engines.image.models.flux.kontext_adapter import ( FluxKontextModelAdapter, ) __all__ = [ "FluxModelAdapter", "FluxKontextModelAdapter", "FLUX_DEV_CONFIG", "FLUX_KONTEXT_CONFIG", "FLUX_SCHNELL_CONFIG", ] ================================================ FILE: src/exo/worker/engines/image/models/flux/adapter.py ================================================ from pathlib import Path from typing import Any import mlx.core as mx from mflux.models.common.config.config import Config from mflux.models.common.config.model_config import ModelConfig from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.model.flux_text_encoder.prompt_encoder import PromptEncoder from mflux.models.flux.model.flux_transformer.transformer import Transformer from mflux.models.flux.variants.txt2img.flux import Flux1 from exo.worker.engines.image.config import ImageModelConfig from exo.worker.engines.image.models.base import ( ModelAdapter, PromptData, RotaryEmbeddings, ) from exo.worker.engines.image.models.flux.wrappers import ( FluxJointBlockWrapper, FluxSingleBlockWrapper, ) from exo.worker.engines.image.pipeline.block_wrapper import ( JointBlockWrapper, SingleBlockWrapper, ) class FluxPromptData(PromptData): def __init__(self, prompt_embeds: mx.array, pooled_prompt_embeds: mx.array): self._prompt_embeds = prompt_embeds self._pooled_prompt_embeds = pooled_prompt_embeds @property def prompt_embeds(self) -> mx.array: return self._prompt_embeds @property def pooled_prompt_embeds(self) -> mx.array: return self._pooled_prompt_embeds @property def negative_prompt_embeds(self) -> mx.array | None: return None @property def negative_pooled_prompt_embeds(self) -> mx.array | None: return None def get_encoder_hidden_states_mask(self, positive: bool = True) -> mx.array | None: return None @property def cond_image_grid( self, ) -> tuple[int, int, int] | list[tuple[int, int, int]] | None: return None @property def conditioning_latents(self) -> mx.array | None: return None @property def kontext_image_ids(self) -> mx.array | None: return None def get_batched_cfg_data( self, ) -> tuple[mx.array, mx.array, mx.array | None, mx.array | None] | None: return None def get_cfg_branch_data( self, positive: bool ) -> tuple[mx.array, mx.array | None, mx.array | None, mx.array | None]: """Flux doesn't use CFG, but we return positive data for compatibility.""" return (self._prompt_embeds, None, self._pooled_prompt_embeds, None) class FluxModelAdapter(ModelAdapter[Flux1, Transformer]): def __init__( self, config: ImageModelConfig, model_id: str, local_path: Path, quantize: int | None = None, ): self._config = config self._model = Flux1( model_config=ModelConfig.from_name(model_name=model_id, base_model=None), model_path=str(local_path), quantize=quantize, ) self._transformer = self._model.transformer @property def hidden_dim(self) -> int: return self._transformer.x_embedder.weight.shape[0] @property def needs_cfg(self) -> bool: return False def _get_latent_creator(self) -> type: return FluxLatentCreator def get_joint_block_wrappers( self, text_seq_len: int, encoder_hidden_states_mask: mx.array | None = None, ) -> list[JointBlockWrapper[Any]]: """Create wrapped joint blocks for Flux.""" return [ FluxJointBlockWrapper(block, text_seq_len) for block in self._transformer.transformer_blocks ] def get_single_block_wrappers( self, text_seq_len: int, ) -> list[SingleBlockWrapper[Any]]: """Create wrapped single blocks for Flux.""" return [ FluxSingleBlockWrapper(block, text_seq_len) for block in self._transformer.single_transformer_blocks ] def slice_transformer_blocks( self, start_layer: int, end_layer: int, ): all_joint = list(self._transformer.transformer_blocks) all_single = list(self._transformer.single_transformer_blocks) total_joint_blocks = len(all_joint) if end_layer <= total_joint_blocks: # All assigned are joint blocks joint_start, joint_end = start_layer, end_layer single_start, single_end = 0, 0 elif start_layer >= total_joint_blocks: # All assigned are single blocks joint_start, joint_end = 0, 0 single_start = start_layer - total_joint_blocks single_end = end_layer - total_joint_blocks else: # Spans both joint and single joint_start, joint_end = start_layer, total_joint_blocks single_start = 0 single_end = end_layer - total_joint_blocks self._transformer.transformer_blocks = all_joint[joint_start:joint_end] self._transformer.single_transformer_blocks = all_single[ single_start:single_end ] def encode_prompt( self, prompt: str, negative_prompt: str | None = None ) -> FluxPromptData: del negative_prompt assert isinstance(self.model.prompt_cache, dict) assert isinstance(self.model.tokenizers, dict) prompt_embeds, pooled_prompt_embeds = PromptEncoder.encode_prompt( prompt=prompt, prompt_cache=self.model.prompt_cache, t5_tokenizer=self.model.tokenizers["t5"], # pyright: ignore[reportAny] clip_tokenizer=self.model.tokenizers["clip"], # pyright: ignore[reportAny] t5_text_encoder=self.model.t5_text_encoder, clip_text_encoder=self.model.clip_text_encoder, ) return FluxPromptData( prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, ) def compute_embeddings( self, hidden_states: mx.array, prompt_embeds: mx.array, ) -> tuple[mx.array, mx.array]: embedded_hidden = self._transformer.x_embedder(hidden_states) embedded_encoder = self._transformer.context_embedder(prompt_embeds) return embedded_hidden, embedded_encoder def compute_text_embeddings( self, t: int, runtime_config: Config, pooled_prompt_embeds: mx.array | None = None, hidden_states: mx.array | None = None, # Ignored by Flux ) -> mx.array: if pooled_prompt_embeds is None: raise ValueError( "pooled_prompt_embeds is required for Flux text embeddings" ) # hidden_states is ignored - Flux uses pooled_prompt_embeds instead return Transformer.compute_text_embeddings( t, pooled_prompt_embeds, self._transformer.time_text_embed, runtime_config ) def compute_rotary_embeddings( self, prompt_embeds: mx.array, runtime_config: Config, encoder_hidden_states_mask: mx.array | None = None, cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]] | None = None, kontext_image_ids: mx.array | None = None, ) -> RotaryEmbeddings: return Transformer.compute_rotary_embeddings( prompt_embeds, self._transformer.pos_embed, runtime_config, kontext_image_ids, ) def apply_guidance( self, noise_positive: mx.array, noise_negative: mx.array, guidance_scale: float, ) -> mx.array: raise NotImplementedError("Flux does not use classifier-free guidance") ================================================ FILE: src/exo/worker/engines/image/models/flux/config.py ================================================ from exo.worker.engines.image.config import ( BlockType, ImageModelConfig, TransformerBlockConfig, ) FLUX_SCHNELL_CONFIG = ImageModelConfig( model_family="flux", block_configs=( TransformerBlockConfig( block_type=BlockType.JOINT, count=19, has_separate_text_output=True ), TransformerBlockConfig( block_type=BlockType.SINGLE, count=38, has_separate_text_output=False ), ), default_steps={"low": 1, "medium": 2, "high": 4}, num_sync_steps=1, ) FLUX_DEV_CONFIG = ImageModelConfig( model_family="flux", block_configs=( TransformerBlockConfig( block_type=BlockType.JOINT, count=19, has_separate_text_output=True ), TransformerBlockConfig( block_type=BlockType.SINGLE, count=38, has_separate_text_output=False ), ), default_steps={"low": 10, "medium": 25, "high": 50}, num_sync_steps=4, ) FLUX_KONTEXT_CONFIG = ImageModelConfig( model_family="flux-kontext", block_configs=( TransformerBlockConfig( block_type=BlockType.JOINT, count=19, has_separate_text_output=True ), TransformerBlockConfig( block_type=BlockType.SINGLE, count=38, has_separate_text_output=False ), ), default_steps={"low": 10, "medium": 25, "high": 50}, num_sync_steps=4, guidance_scale=4.0, ) ================================================ FILE: src/exo/worker/engines/image/models/flux/kontext_adapter.py ================================================ import math from pathlib import Path from typing import Any, final import mlx.core as mx from mflux.models.common.config.config import Config from mflux.models.common.config.model_config import ModelConfig from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.model.flux_text_encoder.prompt_encoder import PromptEncoder from mflux.models.flux.model.flux_transformer.transformer import Transformer from mflux.models.flux.variants.kontext.flux_kontext import Flux1Kontext from mflux.models.flux.variants.kontext.kontext_util import KontextUtil from exo.worker.engines.image.config import ImageModelConfig from exo.worker.engines.image.models.base import ( ModelAdapter, PromptData, RotaryEmbeddings, ) from exo.worker.engines.image.models.flux.wrappers import ( FluxJointBlockWrapper, FluxSingleBlockWrapper, ) from exo.worker.engines.image.pipeline.block_wrapper import ( JointBlockWrapper, SingleBlockWrapper, ) @final class FluxKontextPromptData(PromptData): """Prompt data for FLUX.1-Kontext image editing. Stores text embeddings along with conditioning latents and position IDs for the input image. """ def __init__( self, prompt_embeds: mx.array, pooled_prompt_embeds: mx.array, conditioning_latents: mx.array, kontext_image_ids: mx.array, ): self._prompt_embeds = prompt_embeds self._pooled_prompt_embeds = pooled_prompt_embeds self._conditioning_latents = conditioning_latents self._kontext_image_ids = kontext_image_ids @property def prompt_embeds(self) -> mx.array: return self._prompt_embeds @property def pooled_prompt_embeds(self) -> mx.array: return self._pooled_prompt_embeds @property def negative_prompt_embeds(self) -> mx.array | None: return None @property def negative_pooled_prompt_embeds(self) -> mx.array | None: return None def get_encoder_hidden_states_mask(self, positive: bool = True) -> mx.array | None: return None @property def cond_image_grid( self, ) -> tuple[int, int, int] | list[tuple[int, int, int]] | None: return None @property def conditioning_latents(self) -> mx.array | None: """VAE-encoded input image latents for Kontext conditioning.""" return self._conditioning_latents @property def kontext_image_ids(self) -> mx.array | None: """Position IDs for Kontext conditioning (first_coord=1).""" return self._kontext_image_ids def get_cfg_branch_data( self, positive: bool ) -> tuple[mx.array, mx.array | None, mx.array | None, mx.array | None]: """Kontext doesn't use CFG, but we return positive data for compatibility.""" return ( self._prompt_embeds, None, self._pooled_prompt_embeds, self._conditioning_latents, ) def get_batched_cfg_data( self, ) -> tuple[mx.array, mx.array, mx.array | None, mx.array | None] | None: # Kontext doesn't use CFG return None @final class FluxKontextModelAdapter(ModelAdapter[Flux1Kontext, Transformer]): """Adapter for FLUX.1-Kontext image editing model. Key differences from standard FluxModelAdapter: - Takes an input image and computes output dimensions from it - Creates conditioning latents from the input image via VAE - Creates special position IDs (kontext_image_ids) for conditioning tokens - Creates pure noise latents (not img2img blending) """ def __init__( self, config: ImageModelConfig, model_id: str, local_path: Path, quantize: int | None = None, ): self._config = config self._model = Flux1Kontext( model_config=ModelConfig.from_name(model_name=model_id, base_model=None), model_path=str(local_path), quantize=quantize, ) self._transformer = self._model.transformer # Stores image path and computed dimensions after set_image_dimensions self._image_path: str | None = None self._output_height: int | None = None self._output_width: int | None = None @property def hidden_dim(self) -> int: return self._transformer.x_embedder.weight.shape[0] @property def needs_cfg(self) -> bool: return False def _get_latent_creator(self) -> type: return FluxLatentCreator def get_joint_block_wrappers( self, text_seq_len: int, encoder_hidden_states_mask: mx.array | None = None, ) -> list[JointBlockWrapper[Any]]: """Create wrapped joint blocks for Flux Kontext.""" return [ FluxJointBlockWrapper(block, text_seq_len) for block in self._transformer.transformer_blocks ] def get_single_block_wrappers( self, text_seq_len: int, ) -> list[SingleBlockWrapper[Any]]: """Create wrapped single blocks for Flux Kontext.""" return [ FluxSingleBlockWrapper(block, text_seq_len) for block in self._transformer.single_transformer_blocks ] def slice_transformer_blocks( self, start_layer: int, end_layer: int, ): all_joint = list(self._transformer.transformer_blocks) all_single = list(self._transformer.single_transformer_blocks) total_joint_blocks = len(all_joint) if end_layer <= total_joint_blocks: # All assigned are joint blocks joint_start, joint_end = start_layer, end_layer single_start, single_end = 0, 0 elif start_layer >= total_joint_blocks: # All assigned are single blocks joint_start, joint_end = 0, 0 single_start = start_layer - total_joint_blocks single_end = end_layer - total_joint_blocks else: # Spans both joint and single joint_start, joint_end = start_layer, total_joint_blocks single_start = 0 single_end = end_layer - total_joint_blocks self._transformer.transformer_blocks = all_joint[joint_start:joint_end] self._transformer.single_transformer_blocks = all_single[ single_start:single_end ] def set_image_dimensions(self, image_path: Path) -> tuple[int, int]: """Compute and store dimensions from input image. Also stores image_path for use in encode_prompt(). Args: image_path: Path to the input image Returns: (output_width, output_height) for runtime config """ from mflux.utils.image_util import ImageUtil pil_image = ImageUtil.load_image(str(image_path)).convert("RGB") image_size = pil_image.size # Compute output dimensions from input image aspect ratio # Target area of 1024x1024 = ~1M pixels target_area = 1024 * 1024 ratio = image_size[0] / image_size[1] output_width = math.sqrt(target_area * ratio) output_height = output_width / ratio output_width = round(output_width / 32) * 32 output_height = round(output_height / 32) * 32 # Ensure multiple of 16 for VAE vae_scale_factor = 8 multiple_of = vae_scale_factor * 2 output_width = output_width // multiple_of * multiple_of output_height = output_height // multiple_of * multiple_of self._image_path = str(image_path) self._output_width = int(output_width) self._output_height = int(output_height) return self._output_width, self._output_height def create_latents(self, seed: int, runtime_config: Config) -> mx.array: """Create initial noise latents for Kontext. Unlike standard img2img which blends noise with encoded input, Kontext uses pure noise latents. The input image is provided separately as conditioning. """ return FluxLatentCreator.create_noise( seed=seed, height=runtime_config.height, width=runtime_config.width, ) def encode_prompt( self, prompt: str, negative_prompt: str | None = None ) -> FluxKontextPromptData: """Encode prompt and create conditioning from stored input image. Must call set_image_dimensions() before this method. Args: prompt: Text prompt for editing negative_prompt: Ignored (Kontext doesn't use CFG) Returns: FluxKontextPromptData with text embeddings and image conditioning """ del negative_prompt # Kontext doesn't support negative prompts or CFG if ( self._image_path is None or self._output_height is None or self._output_width is None ): raise RuntimeError( "set_image_dimensions() must be called before encode_prompt() " "for FluxKontextModelAdapter" ) assert isinstance(self.model.prompt_cache, dict) assert isinstance(self.model.tokenizers, dict) # Encode text prompt prompt_embeds, pooled_prompt_embeds = PromptEncoder.encode_prompt( prompt=prompt, prompt_cache=self.model.prompt_cache, t5_tokenizer=self.model.tokenizers["t5"], # pyright: ignore[reportAny] clip_tokenizer=self.model.tokenizers["clip"], # pyright: ignore[reportAny] t5_text_encoder=self.model.t5_text_encoder, clip_text_encoder=self.model.clip_text_encoder, ) # Create conditioning latents from input image conditioning_latents, kontext_image_ids = ( KontextUtil.create_image_conditioning_latents( vae=self.model.vae, height=self._output_height, width=self._output_width, image_path=self._image_path, ) ) return FluxKontextPromptData( prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, conditioning_latents=conditioning_latents, kontext_image_ids=kontext_image_ids, ) def compute_embeddings( self, hidden_states: mx.array, prompt_embeds: mx.array, ) -> tuple[mx.array, mx.array]: embedded_hidden = self._transformer.x_embedder(hidden_states) embedded_encoder = self._transformer.context_embedder(prompt_embeds) return embedded_hidden, embedded_encoder def compute_text_embeddings( self, t: int, runtime_config: Config, pooled_prompt_embeds: mx.array | None = None, hidden_states: mx.array | None = None, ) -> mx.array: if pooled_prompt_embeds is None: raise ValueError( "pooled_prompt_embeds is required for Flux Kontext text embeddings" ) return Transformer.compute_text_embeddings( t, pooled_prompt_embeds, self._transformer.time_text_embed, runtime_config ) def compute_rotary_embeddings( self, prompt_embeds: mx.array, runtime_config: Config, encoder_hidden_states_mask: mx.array | None = None, cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]] | None = None, kontext_image_ids: mx.array | None = None, ) -> RotaryEmbeddings: return Transformer.compute_rotary_embeddings( prompt_embeds, self._transformer.pos_embed, runtime_config, kontext_image_ids, ) def apply_guidance( self, noise_positive: mx.array, noise_negative: mx.array, guidance_scale: float, ) -> mx.array: raise NotImplementedError("Flux Kontext does not use classifier-free guidance") ================================================ FILE: src/exo/worker/engines/image/models/flux/wrappers.py ================================================ from typing import final import mlx.core as mx from mflux.models.flux.model.flux_transformer.common.attention_utils import ( AttentionUtils, ) from mflux.models.flux.model.flux_transformer.joint_transformer_block import ( JointTransformerBlock, ) from mflux.models.flux.model.flux_transformer.single_transformer_block import ( SingleTransformerBlock, ) from pydantic import BaseModel, ConfigDict from exo.worker.engines.image.models.base import RotaryEmbeddings from exo.worker.engines.image.pipeline.block_wrapper import ( JointBlockWrapper, SingleBlockWrapper, ) @final class FluxModulationParams(BaseModel): model_config = ConfigDict(frozen=True, strict=True, arbitrary_types_allowed=True) gate_msa: mx.array shift_mlp: mx.array scale_mlp: mx.array gate_mlp: mx.array @final class FluxNormGateState(BaseModel): model_config = ConfigDict(frozen=True, strict=True, arbitrary_types_allowed=True) norm_hidden: mx.array gate: mx.array class FluxJointBlockWrapper(JointBlockWrapper[JointTransformerBlock]): def __init__(self, block: JointTransformerBlock, text_seq_len: int): super().__init__(block, text_seq_len) self._num_heads = block.attn.num_heads self._head_dim = block.attn.head_dimension # Intermediate state stored between _compute_qkv and _apply_output self._hidden_mod: FluxModulationParams | None = None self._context_mod: FluxModulationParams | None = None def _compute_qkv( self, hidden_states: mx.array, encoder_hidden_states: mx.array, text_embeddings: mx.array, rotary_embeddings: RotaryEmbeddings, patch_mode: bool = False, ) -> tuple[mx.array, mx.array, mx.array]: assert isinstance(rotary_embeddings, mx.array) attn = self.block.attn ( norm_hidden, gate_msa, shift_mlp, scale_mlp, gate_mlp, ) = self.block.norm1( hidden_states=hidden_states, text_embeddings=text_embeddings, ) self._hidden_mod = FluxModulationParams( gate_msa=gate_msa, shift_mlp=shift_mlp, scale_mlp=scale_mlp, gate_mlp=gate_mlp, ) ( norm_encoder, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp, ) = self.block.norm1_context( hidden_states=encoder_hidden_states, text_embeddings=text_embeddings, ) self._context_mod = FluxModulationParams( gate_msa=c_gate_msa, shift_mlp=c_shift_mlp, scale_mlp=c_scale_mlp, gate_mlp=c_gate_mlp, ) img_query, img_key, img_value = AttentionUtils.process_qkv( hidden_states=norm_hidden, to_q=attn.to_q, to_k=attn.to_k, to_v=attn.to_v, norm_q=attn.norm_q, norm_k=attn.norm_k, num_heads=self._num_heads, head_dim=self._head_dim, ) txt_query, txt_key, txt_value = AttentionUtils.process_qkv( hidden_states=norm_encoder, to_q=attn.add_q_proj, to_k=attn.add_k_proj, to_v=attn.add_v_proj, norm_q=attn.norm_added_q, norm_k=attn.norm_added_k, num_heads=self._num_heads, head_dim=self._head_dim, ) query = mx.concatenate([txt_query, img_query], axis=2) key = mx.concatenate([txt_key, img_key], axis=2) value = mx.concatenate([txt_value, img_value], axis=2) if patch_mode: text_rope = rotary_embeddings[:, :, : self._text_seq_len, ...] patch_img_rope = rotary_embeddings[ :, :, self._text_seq_len + self._patch_start : self._text_seq_len + self._patch_end, ..., ] rope = mx.concatenate([text_rope, patch_img_rope], axis=2) else: rope = rotary_embeddings query, key = AttentionUtils.apply_rope(xq=query, xk=key, freqs_cis=rope) return query, key, value def _compute_attention( self, query: mx.array, key: mx.array, value: mx.array ) -> mx.array: batch_size = query.shape[0] return AttentionUtils.compute_attention( query=query, key=key, value=value, batch_size=batch_size, num_heads=self._num_heads, head_dim=self._head_dim, ) def _apply_output( self, attn_out: mx.array, hidden_states: mx.array, encoder_hidden_states: mx.array, text_embeddings: mx.array, ) -> tuple[mx.array, mx.array]: attn = self.block.attn context_attn_output = attn_out[:, : self._text_seq_len, :] hidden_attn_output = attn_out[:, self._text_seq_len :, :] hidden_attn_output = attn.to_out[0](hidden_attn_output) # pyright: ignore[reportAny] context_attn_output = attn.to_add_out(context_attn_output) assert self._hidden_mod is not None assert self._context_mod is not None hidden_states = JointTransformerBlock.apply_norm_and_feed_forward( hidden_states=hidden_states, attn_output=hidden_attn_output, # pyright: ignore[reportAny] gate_mlp=self._hidden_mod.gate_mlp, gate_msa=self._hidden_mod.gate_msa, scale_mlp=self._hidden_mod.scale_mlp, shift_mlp=self._hidden_mod.shift_mlp, norm_layer=self.block.norm2, ff_layer=self.block.ff, ) encoder_hidden_states = JointTransformerBlock.apply_norm_and_feed_forward( hidden_states=encoder_hidden_states, attn_output=context_attn_output, gate_mlp=self._context_mod.gate_mlp, gate_msa=self._context_mod.gate_msa, scale_mlp=self._context_mod.scale_mlp, shift_mlp=self._context_mod.shift_mlp, norm_layer=self.block.norm2_context, ff_layer=self.block.ff_context, ) return encoder_hidden_states, hidden_states class FluxSingleBlockWrapper(SingleBlockWrapper[SingleTransformerBlock]): """Flux-specific single block wrapper with pipefusion support.""" def __init__(self, block: SingleTransformerBlock, text_seq_len: int): super().__init__(block, text_seq_len) self._num_heads = block.attn.num_heads self._head_dim = block.attn.head_dimension # Intermediate state stored between _compute_qkv and _apply_output self._norm_state: FluxNormGateState | None = None def _compute_qkv( self, hidden_states: mx.array, text_embeddings: mx.array, rotary_embeddings: RotaryEmbeddings, patch_mode: bool = False, ) -> tuple[mx.array, mx.array, mx.array]: assert isinstance(rotary_embeddings, mx.array) attn = self.block.attn norm_hidden, gate = self.block.norm( hidden_states=hidden_states, text_embeddings=text_embeddings, ) self._norm_state = FluxNormGateState(norm_hidden=norm_hidden, gate=gate) query, key, value = AttentionUtils.process_qkv( hidden_states=norm_hidden, to_q=attn.to_q, to_k=attn.to_k, to_v=attn.to_v, norm_q=attn.norm_q, norm_k=attn.norm_k, num_heads=self._num_heads, head_dim=self._head_dim, ) if patch_mode: text_rope = rotary_embeddings[:, :, : self._text_seq_len, ...] patch_img_rope = rotary_embeddings[ :, :, self._text_seq_len + self._patch_start : self._text_seq_len + self._patch_end, ..., ] rope = mx.concatenate([text_rope, patch_img_rope], axis=2) else: rope = rotary_embeddings query, key = AttentionUtils.apply_rope(xq=query, xk=key, freqs_cis=rope) return query, key, value def _compute_attention( self, query: mx.array, key: mx.array, value: mx.array ) -> mx.array: batch_size = query.shape[0] return AttentionUtils.compute_attention( query=query, key=key, value=value, batch_size=batch_size, num_heads=self._num_heads, head_dim=self._head_dim, ) def _apply_output( self, attn_out: mx.array, hidden_states: mx.array, text_embeddings: mx.array, ) -> mx.array: residual = hidden_states assert self._norm_state is not None output = self.block._apply_feed_forward_and_projection( norm_hidden_states=self._norm_state.norm_hidden, attn_output=attn_out, gate=self._norm_state.gate, ) return residual + output ================================================ FILE: src/exo/worker/engines/image/models/qwen/__init__.py ================================================ from exo.worker.engines.image.models.qwen.adapter import QwenModelAdapter from exo.worker.engines.image.models.qwen.config import ( QWEN_IMAGE_CONFIG, QWEN_IMAGE_EDIT_CONFIG, ) from exo.worker.engines.image.models.qwen.edit_adapter import QwenEditModelAdapter __all__ = [ "QwenModelAdapter", "QwenEditModelAdapter", "QWEN_IMAGE_CONFIG", "QWEN_IMAGE_EDIT_CONFIG", ] ================================================ FILE: src/exo/worker/engines/image/models/qwen/adapter.py ================================================ from pathlib import Path from typing import Any import mlx.core as mx from mflux.models.common.config import ModelConfig from mflux.models.common.config.config import Config from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator from mflux.models.qwen.model.qwen_text_encoder.qwen_prompt_encoder import ( QwenPromptEncoder, ) from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage from exo.worker.engines.image.config import ImageModelConfig from exo.worker.engines.image.models.base import ( ModelAdapter, PromptData, RotaryEmbeddings, ) from exo.worker.engines.image.models.qwen.wrappers import QwenJointBlockWrapper from exo.worker.engines.image.pipeline.block_wrapper import ( JointBlockWrapper, SingleBlockWrapper, ) class QwenPromptData(PromptData): def __init__( self, prompt_embeds: mx.array, prompt_mask: mx.array, negative_prompt_embeds: mx.array, negative_prompt_mask: mx.array, ): self._prompt_embeds = prompt_embeds self._prompt_mask = prompt_mask self._negative_prompt_embeds = negative_prompt_embeds self._negative_prompt_mask = negative_prompt_mask @property def prompt_embeds(self) -> mx.array: return self._prompt_embeds @property def pooled_prompt_embeds(self) -> mx.array: return self._prompt_embeds @property def negative_prompt_embeds(self) -> mx.array: return self._negative_prompt_embeds @property def negative_pooled_prompt_embeds(self) -> mx.array: return self._negative_prompt_embeds def get_encoder_hidden_states_mask(self, positive: bool = True) -> mx.array: if positive: return self._prompt_mask else: return self._negative_prompt_mask @property def cond_image_grid( self, ) -> tuple[int, int, int] | list[tuple[int, int, int]] | None: return None @property def conditioning_latents(self) -> mx.array | None: return None @property def kontext_image_ids(self) -> mx.array | None: return None def get_batched_cfg_data( self, ) -> tuple[mx.array, mx.array, mx.array | None, mx.array | None] | None: """Batch positive and negative embeddings for CFG with batch_size=2. Pads shorter sequence to max length using zeros for embeddings and zeros (masked) for attention mask. Returns: Tuple of (batched_embeds, batched_mask, None, conditioning_latents) - batched_embeds: [2, max_seq, hidden] - batched_mask: [2, max_seq] - None for pooled (Qwen doesn't use it) - conditioning_latents: [2, latent_seq, latent_dim] or None """ pos_embeds = self._prompt_embeds neg_embeds = self._negative_prompt_embeds pos_mask = self._prompt_mask neg_mask = self._negative_prompt_mask pos_seq_len = pos_embeds.shape[1] neg_seq_len = neg_embeds.shape[1] max_seq_len = max(pos_seq_len, neg_seq_len) hidden_dim = pos_embeds.shape[2] if pos_seq_len < max_seq_len: pad_len = max_seq_len - pos_seq_len pos_embeds = mx.concatenate( [ pos_embeds, mx.zeros((1, pad_len, hidden_dim), dtype=pos_embeds.dtype), ], axis=1, ) pos_mask = mx.concatenate( [pos_mask, mx.zeros((1, pad_len), dtype=pos_mask.dtype)], axis=1, ) elif neg_seq_len < max_seq_len: pad_len = max_seq_len - neg_seq_len neg_embeds = mx.concatenate( [ neg_embeds, mx.zeros((1, pad_len, hidden_dim), dtype=neg_embeds.dtype), ], axis=1, ) neg_mask = mx.concatenate( [neg_mask, mx.zeros((1, pad_len), dtype=neg_mask.dtype)], axis=1, ) batched_embeds = mx.concatenate([pos_embeds, neg_embeds], axis=0) batched_mask = mx.concatenate([pos_mask, neg_mask], axis=0) # TODO(ciaran): currently None but maybe we will deduplicate with edit # adapter cond_latents = self.conditioning_latents if cond_latents is not None: cond_latents = mx.concatenate([cond_latents, cond_latents], axis=0) return batched_embeds, batched_mask, None, cond_latents def get_cfg_branch_data( self, positive: bool ) -> tuple[mx.array, mx.array | None, mx.array | None, mx.array | None]: if positive: return ( self._prompt_embeds, self._prompt_mask, None, self.conditioning_latents, ) else: return ( self._negative_prompt_embeds, self._negative_prompt_mask, None, self.conditioning_latents, ) class QwenModelAdapter(ModelAdapter[QwenImage, QwenTransformer]): """Adapter for Qwen-Image model. Key differences from Flux: - Single text encoder (vs dual T5+CLIP) - 60 joint-style blocks, no single blocks - 3D RoPE returning ((img_cos, img_sin), (txt_cos, txt_sin)) - Norm-preserving CFG with negative prompts - Uses attention mask for variable-length text """ def __init__( self, config: ImageModelConfig, model_id: str, local_path: Path, quantize: int | None = None, ): self._config = config self._model = QwenImage( model_config=ModelConfig.from_name(model_name=model_id, base_model=None), model_path=str(local_path), quantize=quantize, ) self._transformer = self._model.transformer @property def hidden_dim(self) -> int: return self._transformer.inner_dim @property def needs_cfg(self) -> bool: gs = self._config.guidance_scale return gs is not None and gs > 1.0 def _get_latent_creator(self) -> type: return QwenLatentCreator def get_joint_block_wrappers( self, text_seq_len: int, encoder_hidden_states_mask: mx.array | None = None, ) -> list[JointBlockWrapper[Any]]: """Create wrapped joint blocks for Qwen.""" return [ QwenJointBlockWrapper(block, text_seq_len, encoder_hidden_states_mask) for block in self._transformer.transformer_blocks ] def get_single_block_wrappers( self, text_seq_len: int, ) -> list[SingleBlockWrapper[Any]]: return [] def slice_transformer_blocks( self, start_layer: int, end_layer: int, ): self._transformer.transformer_blocks = self._transformer.transformer_blocks[ start_layer:end_layer ] def encode_prompt( self, prompt: str, negative_prompt: str | None = None ) -> QwenPromptData: assert isinstance(self.model.prompt_cache, dict) assert isinstance(self.model.tokenizers, dict) if negative_prompt is None or negative_prompt == "": negative_prompt = " " prompt_embeds, prompt_mask, neg_embeds, neg_mask = ( QwenPromptEncoder.encode_prompt( prompt=prompt, negative_prompt=negative_prompt, prompt_cache=self.model.prompt_cache, qwen_tokenizer=self.model.tokenizers["qwen"], # pyright: ignore[reportAny] qwen_text_encoder=self.model.text_encoder, ) ) return QwenPromptData( prompt_embeds=prompt_embeds, prompt_mask=prompt_mask, negative_prompt_embeds=neg_embeds, negative_prompt_mask=neg_mask, ) def compute_embeddings( self, hidden_states: mx.array, prompt_embeds: mx.array, ) -> tuple[mx.array, mx.array]: embedded_hidden = self._transformer.img_in(hidden_states) encoder_hidden_states = self._transformer.txt_norm(prompt_embeds) embedded_encoder = self._transformer.txt_in(encoder_hidden_states) return embedded_hidden, embedded_encoder def compute_text_embeddings( self, t: int, runtime_config: Config, pooled_prompt_embeds: mx.array | None = None, hidden_states: mx.array | None = None, ) -> mx.array: # Use hidden_states if provided, otherwise fall back to pooled_prompt_embeds # (which for Qwen is the same as prompt_embeds) ref_tensor = ( hidden_states if hidden_states is not None else pooled_prompt_embeds ) if ref_tensor is None: raise ValueError( "Either hidden_states or pooled_prompt_embeds is required " "for Qwen text embeddings" ) timestep = QwenTransformer._compute_timestep(t, runtime_config) # noqa: SLF001 batch_size = ref_tensor.shape[0] timestep = mx.broadcast_to(timestep, (batch_size,)).astype(mx.float32) return self._transformer.time_text_embed(timestep, ref_tensor) # pyright: ignore[reportAny] def compute_rotary_embeddings( self, prompt_embeds: mx.array, runtime_config: Config, encoder_hidden_states_mask: mx.array | None = None, cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]] | None = None, kontext_image_ids: mx.array | None = None, ) -> RotaryEmbeddings: if encoder_hidden_states_mask is None: raise ValueError( "encoder_hidden_states_mask is required for Qwen RoPE computation" ) return QwenTransformer._compute_rotary_embeddings( encoder_hidden_states_mask=encoder_hidden_states_mask, pos_embed=self._transformer.pos_embed, # pyright: ignore[reportAny] config=runtime_config, cond_image_grid=cond_image_grid, ) def apply_guidance( self, noise_positive: mx.array, noise_negative: mx.array, guidance_scale: float, ) -> mx.array: return self._model.compute_guided_noise( noise=noise_positive, noise_negative=noise_negative, guidance=guidance_scale, ) ================================================ FILE: src/exo/worker/engines/image/models/qwen/config.py ================================================ from exo.worker.engines.image.config import ( BlockType, ImageModelConfig, TransformerBlockConfig, ) QWEN_IMAGE_CONFIG = ImageModelConfig( model_family="qwen", block_configs=( TransformerBlockConfig( block_type=BlockType.JOINT, count=60, has_separate_text_output=True ), ), default_steps={"low": 10, "medium": 25, "high": 50}, num_sync_steps=7, guidance_scale=3.5, # Set to None or < 1.0 to disable CFG ) QWEN_IMAGE_EDIT_CONFIG = ImageModelConfig( model_family="qwen-edit", block_configs=( TransformerBlockConfig( block_type=BlockType.JOINT, count=60, has_separate_text_output=True ), ), default_steps={"low": 10, "medium": 25, "high": 50}, num_sync_steps=7, guidance_scale=3.5, ) ================================================ FILE: src/exo/worker/engines/image/models/qwen/edit_adapter.py ================================================ import math from dataclasses import dataclass from pathlib import Path from typing import Any import mlx.core as mx from mflux.models.common.config.config import Config from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer from mflux.models.qwen.variants.edit.qwen_edit_util import QwenEditUtil from mflux.models.qwen.variants.edit.qwen_image_edit import QwenImageEdit from exo.worker.engines.image.config import ImageModelConfig from exo.worker.engines.image.models.base import ( ModelAdapter, PromptData, RotaryEmbeddings, ) from exo.worker.engines.image.models.qwen.wrappers import QwenJointBlockWrapper from exo.worker.engines.image.pipeline.block_wrapper import ( JointBlockWrapper, SingleBlockWrapper, ) @dataclass(frozen=True) class EditImageDimensions: vl_width: int vl_height: int vae_width: int vae_height: int image_paths: list[str] class QwenEditPromptData(PromptData): def __init__( self, prompt_embeds: mx.array, prompt_mask: mx.array, negative_prompt_embeds: mx.array, negative_prompt_mask: mx.array, conditioning_latents: mx.array, qwen_image_ids: mx.array, cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]], ): self._prompt_embeds = prompt_embeds self._prompt_mask = prompt_mask self._negative_prompt_embeds = negative_prompt_embeds self._negative_prompt_mask = negative_prompt_mask self._conditioning_latents = conditioning_latents self._qwen_image_ids = qwen_image_ids self._cond_image_grid = cond_image_grid @property def prompt_embeds(self) -> mx.array: return self._prompt_embeds @property def pooled_prompt_embeds(self) -> mx.array: return self._prompt_embeds @property def negative_prompt_embeds(self) -> mx.array: return self._negative_prompt_embeds @property def negative_pooled_prompt_embeds(self) -> mx.array: return self._negative_prompt_embeds def get_encoder_hidden_states_mask(self, positive: bool = True) -> mx.array: if positive: return self._prompt_mask else: return self._negative_prompt_mask @property def cond_image_grid(self) -> tuple[int, int, int] | list[tuple[int, int, int]]: return self._cond_image_grid @property def conditioning_latents(self) -> mx.array: return self._conditioning_latents @property def qwen_image_ids(self) -> mx.array: return self._qwen_image_ids @property def kontext_image_ids(self) -> mx.array | None: return None @property def is_edit_mode(self) -> bool: return True def get_batched_cfg_data( self, ) -> tuple[mx.array, mx.array, mx.array | None, mx.array | None] | None: """Batch positive and negative embeddings for CFG with batch_size=2. Pads shorter sequence to max length using zeros for embeddings and zeros (masked) for attention mask. Duplicates conditioning latents for both positive and negative passes. Returns: Tuple of (batched_embeds, batched_mask, None, batched_cond_latents) - batched_embeds: [2, max_seq, hidden] - batched_mask: [2, max_seq] - None for pooled (Qwen doesn't use it) - batched_cond_latents: [2, latent_seq, latent_dim] TODO(ciaran): type this """ pos_embeds = self._prompt_embeds neg_embeds = self._negative_prompt_embeds pos_mask = self._prompt_mask neg_mask = self._negative_prompt_mask pos_seq_len = pos_embeds.shape[1] neg_seq_len = neg_embeds.shape[1] max_seq_len = max(pos_seq_len, neg_seq_len) hidden_dim = pos_embeds.shape[2] if pos_seq_len < max_seq_len: pad_len = max_seq_len - pos_seq_len pos_embeds = mx.concatenate( [ pos_embeds, mx.zeros((1, pad_len, hidden_dim), dtype=pos_embeds.dtype), ], axis=1, ) pos_mask = mx.concatenate( [pos_mask, mx.zeros((1, pad_len), dtype=pos_mask.dtype)], axis=1, ) if neg_seq_len < max_seq_len: pad_len = max_seq_len - neg_seq_len neg_embeds = mx.concatenate( [ neg_embeds, mx.zeros((1, pad_len, hidden_dim), dtype=neg_embeds.dtype), ], axis=1, ) neg_mask = mx.concatenate( [neg_mask, mx.zeros((1, pad_len), dtype=neg_mask.dtype)], axis=1, ) batched_embeds = mx.concatenate([pos_embeds, neg_embeds], axis=0) batched_mask = mx.concatenate([pos_mask, neg_mask], axis=0) batched_cond_latents = mx.concatenate( [self._conditioning_latents, self._conditioning_latents], axis=0 ) return batched_embeds, batched_mask, None, batched_cond_latents def get_cfg_branch_data( self, positive: bool ) -> tuple[mx.array, mx.array | None, mx.array | None, mx.array | None]: if positive: return ( self._prompt_embeds, self._prompt_mask, None, self._conditioning_latents, ) else: return ( self._negative_prompt_embeds, self._negative_prompt_mask, None, self._conditioning_latents, ) class QwenEditModelAdapter(ModelAdapter[QwenImageEdit, QwenTransformer]): """Adapter for Qwen-Image-Edit model. Key differences from standard QwenModelAdapter: - Uses QwenImageEdit model with vision-language components - Encodes prompts WITH input images via VL tokenizer/encoder - Creates conditioning latents from input images - Supports image editing with concatenated latents during diffusion """ def __init__( self, config: ImageModelConfig, model_id: str, local_path: Path, quantize: int | None = None, ): self._config = config self._model = QwenImageEdit( quantize=quantize, model_path=str(local_path), ) self._transformer = self._model.transformer self._edit_dimensions: EditImageDimensions | None = None @property def config(self) -> ImageModelConfig: return self._config @property def model(self) -> QwenImageEdit: return self._model @property def transformer(self) -> QwenTransformer: return self._transformer @property def hidden_dim(self) -> int: return self._transformer.inner_dim @property def needs_cfg(self) -> bool: gs = self._config.guidance_scale return gs is not None and gs > 1.0 def _get_latent_creator(self) -> type[QwenLatentCreator]: return QwenLatentCreator def get_joint_block_wrappers( self, text_seq_len: int, encoder_hidden_states_mask: mx.array | None = None, ) -> list[JointBlockWrapper[Any]]: """Create wrapped joint blocks for Qwen Edit.""" return [ QwenJointBlockWrapper(block, text_seq_len, encoder_hidden_states_mask) for block in self._transformer.transformer_blocks ] def get_single_block_wrappers( self, text_seq_len: int, ) -> list[SingleBlockWrapper[Any]]: """Qwen has no single blocks.""" return [] def slice_transformer_blocks( self, start_layer: int, end_layer: int, ): self._transformer.transformer_blocks = self._transformer.transformer_blocks[ start_layer:end_layer ] def set_image_dimensions(self, image_path: Path) -> tuple[int, int]: """Compute and store dimensions from input image. Also stores image_paths for use in encode_prompt(). Returns: (output_width, output_height) for runtime config """ vl_w, vl_h, vae_w, vae_h, out_w, out_h = self._compute_dimensions_from_image( image_path ) self._edit_dimensions = EditImageDimensions( vl_width=vl_w, vl_height=vl_h, vae_width=vae_w, vae_height=vae_h, image_paths=[str(image_path)], ) return out_w, out_h def create_latents(self, seed: int, runtime_config: Config) -> mx.array: """Create initial noise latents (pure noise for edit mode).""" return QwenLatentCreator.create_noise( seed=seed, height=runtime_config.height, width=runtime_config.width, ) def encode_prompt( self, prompt: str, negative_prompt: str | None = None ) -> QwenEditPromptData: dims = self._edit_dimensions if dims is None: raise RuntimeError( "set_image_dimensions() must be called before encode_prompt() " "for QwenEditModelAdapter" ) if negative_prompt is None or negative_prompt == "": negative_prompt = " " # TODO(ciaran): config is untyped and unused, unsure if Config or RuntimeConfig is intended ( prompt_embeds, prompt_mask, negative_prompt_embeds, negative_prompt_mask, ) = self._model._encode_prompts_with_images( prompt, negative_prompt, dims.image_paths, self._config, # pyright: ignore[reportArgumentType] dims.vl_width, dims.vl_height, ) ( conditioning_latents, qwen_image_ids, cond_h_patches, cond_w_patches, num_images, ) = QwenEditUtil.create_image_conditioning_latents( # pyright: ignore[reportUnknownMemberType] vae=self._model.vae, height=dims.vae_height, width=dims.vae_width, image_paths=dims.image_paths, vl_width=dims.vl_width, vl_height=dims.vl_height, ) if num_images > 1: cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]] = [ (1, cond_h_patches, cond_w_patches) for _ in range(num_images) ] else: cond_image_grid = (1, cond_h_patches, cond_w_patches) return QwenEditPromptData( prompt_embeds=prompt_embeds, prompt_mask=prompt_mask, negative_prompt_embeds=negative_prompt_embeds, negative_prompt_mask=negative_prompt_mask, conditioning_latents=conditioning_latents, qwen_image_ids=qwen_image_ids, cond_image_grid=cond_image_grid, ) def compute_embeddings( self, hidden_states: mx.array, prompt_embeds: mx.array, ) -> tuple[mx.array, mx.array]: embedded_hidden = self._transformer.img_in(hidden_states) encoder_hidden_states = self._transformer.txt_norm(prompt_embeds) embedded_encoder = self._transformer.txt_in(encoder_hidden_states) return embedded_hidden, embedded_encoder def compute_text_embeddings( self, t: int, runtime_config: Config, pooled_prompt_embeds: mx.array | None = None, hidden_states: mx.array | None = None, ) -> mx.array: ref_tensor = ( hidden_states if hidden_states is not None else pooled_prompt_embeds ) if ref_tensor is None: raise ValueError( "Either hidden_states or pooled_prompt_embeds is required " "for Qwen text embeddings" ) timestep = QwenTransformer._compute_timestep(t, runtime_config) # noqa: SLF001 batch_size = ref_tensor.shape[0] timestep = mx.broadcast_to(timestep, (batch_size,)).astype(mx.float32) return self._transformer.time_text_embed(timestep, ref_tensor) # pyright: ignore[reportAny] def compute_rotary_embeddings( self, prompt_embeds: mx.array, runtime_config: Config, encoder_hidden_states_mask: mx.array | None = None, cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]] | None = None, kontext_image_ids: mx.array | None = None, ) -> RotaryEmbeddings: if encoder_hidden_states_mask is None: raise ValueError( "encoder_hidden_states_mask is required for Qwen RoPE computation" ) return QwenTransformer._compute_rotary_embeddings( encoder_hidden_states_mask=encoder_hidden_states_mask, pos_embed=self._transformer.pos_embed, # pyright: ignore[reportAny] config=runtime_config, cond_image_grid=cond_image_grid, ) def apply_guidance( self, noise_positive: mx.array, noise_negative: mx.array, guidance_scale: float, ) -> mx.array: from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage return QwenImage.compute_guided_noise( noise=noise_positive, noise_negative=noise_negative, guidance=guidance_scale, ) def _compute_dimensions_from_image( self, image_path: Path ) -> tuple[int, int, int, int, int, int]: from mflux.utils.image_util import ImageUtil pil_image = ImageUtil.load_image(str(image_path)).convert("RGB") image_size = pil_image.size # Vision-language dimensions (384x384 target area) condition_image_size = 384 * 384 condition_ratio = image_size[0] / image_size[1] vl_width = math.sqrt(condition_image_size * condition_ratio) vl_height = vl_width / condition_ratio vl_width = round(vl_width / 32) * 32 vl_height = round(vl_height / 32) * 32 # VAE dimensions (1024x1024 target area) vae_image_size = 1024 * 1024 vae_ratio = image_size[0] / image_size[1] vae_width = math.sqrt(vae_image_size * vae_ratio) vae_height = vae_width / vae_ratio vae_width = round(vae_width / 32) * 32 vae_height = round(vae_height / 32) * 32 # Output dimensions from input image aspect ratio target_area = 1024 * 1024 ratio = image_size[0] / image_size[1] output_width = math.sqrt(target_area * ratio) output_height = output_width / ratio output_width = round(output_width / 32) * 32 output_height = round(output_height / 32) * 32 # Ensure multiple of 16 for VAE vae_scale_factor = 8 multiple_of = vae_scale_factor * 2 output_width = output_width // multiple_of * multiple_of output_height = output_height // multiple_of * multiple_of return ( int(vl_width), int(vl_height), int(vae_width), int(vae_height), int(output_width), int(output_height), ) ================================================ FILE: src/exo/worker/engines/image/models/qwen/wrappers.py ================================================ from typing import final import mlx.core as mx from mflux.models.qwen.model.qwen_transformer.qwen_attention import QwenAttention from mflux.models.qwen.model.qwen_transformer.qwen_transformer_block import ( QwenTransformerBlock, ) from pydantic import BaseModel, ConfigDict from exo.worker.engines.image.models.base import RotaryEmbeddings from exo.worker.engines.image.pipeline.block_wrapper import JointBlockWrapper @final class QwenStreamModulation(BaseModel): model_config = ConfigDict(frozen=True, strict=True, arbitrary_types_allowed=True) mod1: mx.array mod2: mx.array gate1: mx.array class QwenJointBlockWrapper(JointBlockWrapper[QwenTransformerBlock]): def __init__( self, block: QwenTransformerBlock, text_seq_len: int, encoder_hidden_states_mask: mx.array | None = None, ): super().__init__(block, text_seq_len) self._encoder_hidden_states_mask = encoder_hidden_states_mask self._num_heads = block.attn.num_heads self._head_dim = block.attn.head_dim # Intermediate state stored between _compute_qkv and _apply_output self._img_mod: QwenStreamModulation | None = None self._txt_mod: QwenStreamModulation | None = None def set_encoder_mask(self, mask: mx.array | None) -> None: """Set the encoder hidden states mask for attention.""" self._encoder_hidden_states_mask = mask def _compute_qkv( self, hidden_states: mx.array, encoder_hidden_states: mx.array, text_embeddings: mx.array, rotary_embeddings: RotaryEmbeddings, patch_mode: bool = False, ) -> tuple[mx.array, mx.array, mx.array]: assert isinstance(rotary_embeddings, tuple) batch_size = hidden_states.shape[0] img_seq_len = hidden_states.shape[1] attn = self.block.attn img_mod_params = self.block.img_mod_linear( self.block.img_mod_silu(text_embeddings) # pyright: ignore[reportUnknownArgumentType] ) txt_mod_params = self.block.txt_mod_linear( self.block.txt_mod_silu(text_embeddings) # pyright: ignore[reportUnknownArgumentType] ) img_mod1, img_mod2 = mx.split(img_mod_params, 2, axis=-1) txt_mod1, txt_mod2 = mx.split(txt_mod_params, 2, axis=-1) img_normed = self.block.img_norm1(hidden_states) img_modulated, img_gate1 = QwenTransformerBlock._modulate(img_normed, img_mod1) self._img_mod = QwenStreamModulation( mod1=img_mod1, mod2=img_mod2, gate1=img_gate1 ) txt_normed = self.block.txt_norm1(encoder_hidden_states) txt_modulated, txt_gate1 = QwenTransformerBlock._modulate(txt_normed, txt_mod1) self._txt_mod = QwenStreamModulation( mod1=txt_mod1, mod2=txt_mod2, gate1=txt_gate1 ) img_query = attn.to_q(img_modulated) img_key = attn.to_k(img_modulated) img_value = attn.to_v(img_modulated) txt_query = attn.add_q_proj(txt_modulated) txt_key = attn.add_k_proj(txt_modulated) txt_value = attn.add_v_proj(txt_modulated) img_query = mx.reshape( img_query, (batch_size, img_seq_len, self._num_heads, self._head_dim) ) img_key = mx.reshape( img_key, (batch_size, img_seq_len, self._num_heads, self._head_dim) ) img_value = mx.reshape( img_value, (batch_size, img_seq_len, self._num_heads, self._head_dim) ) txt_query = mx.reshape( txt_query, (batch_size, self._text_seq_len, self._num_heads, self._head_dim), ) txt_key = mx.reshape( txt_key, (batch_size, self._text_seq_len, self._num_heads, self._head_dim) ) txt_value = mx.reshape( txt_value, (batch_size, self._text_seq_len, self._num_heads, self._head_dim) ) img_query = attn.norm_q(img_query) img_key = attn.norm_k(img_key) txt_query = attn.norm_added_q(txt_query) txt_key = attn.norm_added_k(txt_key) (img_cos, img_sin), (txt_cos, txt_sin) = rotary_embeddings if patch_mode: # Slice image RoPE for patch, keep full text RoPE img_cos = img_cos[self._patch_start : self._patch_end] img_sin = img_sin[self._patch_start : self._patch_end] img_query = QwenAttention._apply_rope_qwen(img_query, img_cos, img_sin) img_key = QwenAttention._apply_rope_qwen(img_key, img_cos, img_sin) txt_query = QwenAttention._apply_rope_qwen(txt_query, txt_cos, txt_sin) txt_key = QwenAttention._apply_rope_qwen(txt_key, txt_cos, txt_sin) img_query = mx.transpose(img_query, (0, 2, 1, 3)) img_key = mx.transpose(img_key, (0, 2, 1, 3)) img_value = mx.transpose(img_value, (0, 2, 1, 3)) txt_query = mx.transpose(txt_query, (0, 2, 1, 3)) txt_key = mx.transpose(txt_key, (0, 2, 1, 3)) txt_value = mx.transpose(txt_value, (0, 2, 1, 3)) query = mx.concatenate([txt_query, img_query], axis=2) key = mx.concatenate([txt_key, img_key], axis=2) value = mx.concatenate([txt_value, img_value], axis=2) return query, key, value def _compute_attention( self, query: mx.array, key: mx.array, value: mx.array ) -> mx.array: attn = self.block.attn mask = QwenAttention._convert_mask_for_qwen( mask=self._encoder_hidden_states_mask, joint_seq_len=key.shape[2], txt_seq_len=self._text_seq_len, ) query_bshd = mx.transpose(query, (0, 2, 1, 3)) key_bshd = mx.transpose(key, (0, 2, 1, 3)) value_bshd = mx.transpose(value, (0, 2, 1, 3)) return attn._compute_attention_qwen( query=query_bshd, key=key_bshd, value=value_bshd, mask=mask, block_idx=None, ) def _apply_output( self, attn_out: mx.array, hidden_states: mx.array, encoder_hidden_states: mx.array, text_embeddings: mx.array, ) -> tuple[mx.array, mx.array]: attn = self.block.attn assert self._img_mod is not None assert self._txt_mod is not None txt_attn_output = attn_out[:, : self._text_seq_len, :] img_attn_output = attn_out[:, self._text_seq_len :, :] img_attn_output = attn.attn_to_out[0](img_attn_output) # pyright: ignore[reportAny] txt_attn_output = attn.to_add_out(txt_attn_output) hidden_states = hidden_states + self._img_mod.gate1 * img_attn_output # pyright: ignore[reportAny] encoder_hidden_states = ( encoder_hidden_states + self._txt_mod.gate1 * txt_attn_output ) img_normed2 = self.block.img_norm2(hidden_states) img_modulated2, img_gate2 = QwenTransformerBlock._modulate( img_normed2, self._img_mod.mod2 ) img_mlp_output = self.block.img_ff(img_modulated2) # pyright: ignore[reportAny] hidden_states = hidden_states + img_gate2 * img_mlp_output # pyright: ignore[reportAny] txt_normed2 = self.block.txt_norm2(encoder_hidden_states) txt_modulated2, txt_gate2 = QwenTransformerBlock._modulate( txt_normed2, self._txt_mod.mod2 ) txt_mlp_output = self.block.txt_ff(txt_modulated2) # pyright: ignore[reportAny] encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output # pyright: ignore[reportAny] return encoder_hidden_states, hidden_states ================================================ FILE: src/exo/worker/engines/image/pipeline/__init__.py ================================================ from exo.worker.engines.image.pipeline.block_wrapper import ( BlockWrapperMode, JointBlockWrapper, SingleBlockWrapper, ) from exo.worker.engines.image.pipeline.kv_cache import ImagePatchKVCache from exo.worker.engines.image.pipeline.runner import DiffusionRunner __all__ = [ "BlockWrapperMode", "DiffusionRunner", "ImagePatchKVCache", "JointBlockWrapper", "SingleBlockWrapper", ] ================================================ FILE: src/exo/worker/engines/image/pipeline/block_wrapper.py ================================================ from abc import ABC, abstractmethod from enum import Enum from typing import Generic, Self, TypeVar import mlx.core as mx from exo.worker.engines.image.models.base import RotaryEmbeddings from exo.worker.engines.image.pipeline.kv_cache import ImagePatchKVCache BlockT = TypeVar("BlockT") class BlockWrapperMode(Enum): CACHING = "caching" # Sync mode: compute full attention, populate cache PATCHED = "patched" # Async mode: compute patch attention, use cached KV class BlockWrapperMixin: """Common cache management logic for block wrappers. Including: - KV cache creation and management - Mode - Patch range setting """ _text_seq_len: int _kv_cache: ImagePatchKVCache | None _mode: BlockWrapperMode _patch_start: int _patch_end: int def _init_cache_state(self, text_seq_len: int) -> None: self._text_seq_len = text_seq_len self._kv_cache = None self._mode = BlockWrapperMode.CACHING self._patch_start = 0 self._patch_end = 0 def set_patch( self, mode: BlockWrapperMode, patch_start: int = 0, patch_end: int = 0, ) -> Self: """Set mode and patch range. Args: mode: CACHING (full attention) or PATCHED (use cached KV) patch_start: Start token index within image (for PATCHED mode) patch_end: End token index within image (for PATCHED mode) Returns: Self for method chaining """ self._mode = mode self._patch_start = patch_start self._patch_end = patch_end return self def set_text_seq_len(self, text_seq_len: int) -> None: self._text_seq_len = text_seq_len def _get_active_cache(self) -> ImagePatchKVCache | None: return self._kv_cache def _ensure_cache(self, img_key: mx.array) -> None: if self._kv_cache is None: batch, num_heads, img_seq_len, head_dim = img_key.shape self._kv_cache = ImagePatchKVCache( batch_size=batch, num_heads=num_heads, image_seq_len=img_seq_len, head_dim=head_dim, ) def _cache_full_image_kv(self, img_key: mx.array, img_value: mx.array) -> None: self._ensure_cache(img_key) cache = self._get_active_cache() assert cache is not None cache.update_image_patch(0, img_key.shape[2], img_key, img_value) def _cache_patch_kv(self, img_key: mx.array, img_value: mx.array) -> None: cache = self._get_active_cache() assert cache is not None cache.update_image_patch(self._patch_start, self._patch_end, img_key, img_value) def _get_full_kv( self, text_key: mx.array, text_value: mx.array ) -> tuple[mx.array, mx.array]: cache = self._get_active_cache() assert cache is not None return cache.get_full_kv(text_key, text_value) def reset_cache(self) -> None: self._kv_cache = None class JointBlockWrapper(BlockWrapperMixin, ABC, Generic[BlockT]): """Base class for joint transformer block wrappers with pipefusion support. The wrapper: - Owns its KV cache (created lazily on first CACHING forward) - Controls the forward pass flow (CACHING vs PATCHED mode) - Handles patch slicing and cache operations """ block: BlockT def __init__(self, block: BlockT, text_seq_len: int): self.block = block self._init_cache_state(text_seq_len) def set_encoder_mask(self, mask: mx.array | None) -> None: # noqa: B027 """Set the encoder hidden states mask for attention. Override in subclasses that use attention masks Default is a no-op for models that don't use masks """ del mask # Unused in base class def __call__( self, hidden_states: mx.array, encoder_hidden_states: mx.array, text_embeddings: mx.array, rotary_embeddings: RotaryEmbeddings, ) -> tuple[mx.array, mx.array]: if self._mode == BlockWrapperMode.CACHING: return self._forward_caching( hidden_states, encoder_hidden_states, text_embeddings, rotary_embeddings ) return self._forward_patched( hidden_states, encoder_hidden_states, text_embeddings, rotary_embeddings ) def _forward_caching( self, hidden_states: mx.array, encoder_hidden_states: mx.array, text_embeddings: mx.array, rotary_embeddings: RotaryEmbeddings, ) -> tuple[mx.array, mx.array]: """CACHING mode: Full attention, store image K/V in cache.""" query, key, value = self._compute_qkv( hidden_states, encoder_hidden_states, text_embeddings, rotary_embeddings ) img_key = key[:, :, self._text_seq_len :, :] img_value = value[:, :, self._text_seq_len :, :] self._cache_full_image_kv(img_key, img_value) attn_out = self._compute_attention(query, key, value) return self._apply_output( attn_out, hidden_states, encoder_hidden_states, text_embeddings ) def _forward_patched( self, hidden_states: mx.array, encoder_hidden_states: mx.array, text_embeddings: mx.array, rotary_embeddings: RotaryEmbeddings, ) -> tuple[mx.array, mx.array]: # hidden_states is already the patch (provided by runner) patch_hidden = hidden_states query, key, value = self._compute_qkv( patch_hidden, encoder_hidden_states, text_embeddings, rotary_embeddings, patch_mode=True, ) text_key = key[:, :, : self._text_seq_len, :] text_value = value[:, :, : self._text_seq_len, :] img_key = key[:, :, self._text_seq_len :, :] img_value = value[:, :, self._text_seq_len :, :] self._cache_patch_kv(img_key, img_value) full_key, full_value = self._get_full_kv(text_key, text_value) attn_out = self._compute_attention(query, full_key, full_value) return self._apply_output( attn_out, patch_hidden, encoder_hidden_states, text_embeddings ) @abstractmethod def _compute_qkv( self, hidden_states: mx.array, encoder_hidden_states: mx.array, text_embeddings: mx.array, rotary_embeddings: RotaryEmbeddings, patch_mode: bool = False, ) -> tuple[mx.array, mx.array, mx.array]: ... @abstractmethod def _compute_attention( self, query: mx.array, key: mx.array, value: mx.array ) -> mx.array: ... @abstractmethod def _apply_output( self, attn_out: mx.array, hidden_states: mx.array, encoder_hidden_states: mx.array, text_embeddings: mx.array, ) -> tuple[mx.array, mx.array]: ... class SingleBlockWrapper(BlockWrapperMixin, ABC, Generic[BlockT]): """Base class for single-stream transformer block wrappers. Similar to JointBlockWrapper but for blocks that operate on a single concatenated [text, image] stream rather than separate streams. """ block: BlockT def __init__(self, block: BlockT, text_seq_len: int): self.block = block self._init_cache_state(text_seq_len) def __call__( self, hidden_states: mx.array, text_embeddings: mx.array, rotary_embeddings: RotaryEmbeddings, ) -> mx.array: if self._mode == BlockWrapperMode.CACHING: return self._forward_caching( hidden_states, text_embeddings, rotary_embeddings ) return self._forward_patched(hidden_states, text_embeddings, rotary_embeddings) def _forward_caching( self, hidden_states: mx.array, text_embeddings: mx.array, rotary_embeddings: RotaryEmbeddings, ) -> mx.array: """CACHING mode: Full attention, store image K/V in cache.""" query, key, value = self._compute_qkv( hidden_states, text_embeddings, rotary_embeddings ) img_key = key[:, :, self._text_seq_len :, :] img_value = value[:, :, self._text_seq_len :, :] self._cache_full_image_kv(img_key, img_value) attn_out = self._compute_attention(query, key, value) return self._apply_output(attn_out, hidden_states, text_embeddings) def _forward_patched( self, hidden_states: mx.array, text_embeddings: mx.array, rotary_embeddings: RotaryEmbeddings, ) -> mx.array: """PATCHED mode: Compute patch Q/K/V, use cached image K/V for attention.""" query, key, value = self._compute_qkv( hidden_states, text_embeddings, rotary_embeddings, patch_mode=True ) text_key = key[:, :, : self._text_seq_len, :] text_value = value[:, :, : self._text_seq_len, :] img_key = key[:, :, self._text_seq_len :, :] img_value = value[:, :, self._text_seq_len :, :] self._cache_patch_kv(img_key, img_value) full_key, full_value = self._get_full_kv(text_key, text_value) attn_out = self._compute_attention(query, full_key, full_value) return self._apply_output(attn_out, hidden_states, text_embeddings) @abstractmethod def _compute_qkv( self, hidden_states: mx.array, text_embeddings: mx.array, rotary_embeddings: RotaryEmbeddings, patch_mode: bool = False, ) -> tuple[mx.array, mx.array, mx.array]: ... @abstractmethod def _compute_attention( self, query: mx.array, key: mx.array, value: mx.array ) -> mx.array: ... @abstractmethod def _apply_output( self, attn_out: mx.array, hidden_states: mx.array, text_embeddings: mx.array, ) -> mx.array: ... ================================================ FILE: src/exo/worker/engines/image/pipeline/kv_cache.py ================================================ import mlx.core as mx class ImagePatchKVCache: """KV cache that stores only IMAGE K/V with patch-level updates. Only caches image K/V since: - Text K/V is always computed fresh (same for all patches) - Only image portion needs stale/fresh cache management across patches """ def __init__( self, batch_size: int, num_heads: int, image_seq_len: int, head_dim: int, dtype: mx.Dtype = mx.float32, ): self.batch_size = batch_size self.num_heads = num_heads self.image_seq_len = image_seq_len self.head_dim = head_dim self._dtype = dtype self.key_cache = mx.zeros( (batch_size, num_heads, image_seq_len, head_dim), dtype=dtype ) self.value_cache = mx.zeros( (batch_size, num_heads, image_seq_len, head_dim), dtype=dtype ) def update_image_patch( self, patch_start: int, patch_end: int, key: mx.array, value: mx.array ) -> None: """Update cache with fresh K/V for an image patch slice. Args: patch_start: Start token index within image portion (0-indexed) patch_end: End token index within image portion key: Fresh key tensor [batch, heads, patch_seq_len, head_dim] value: Fresh value tensor [batch, heads, patch_seq_len, head_dim] """ self.key_cache[:, :, patch_start:patch_end, :] = key self.value_cache[:, :, patch_start:patch_end, :] = value def get_full_kv( self, text_key: mx.array, text_value: mx.array ) -> tuple[mx.array, mx.array]: """Return full K/V by concatenating fresh text K/V with cached image K/V. Args: text_key: Fresh text key tensor [batch, heads, text_seq_len, head_dim] text_value: Fresh text value tensor [batch, heads, text_seq_len, head_dim] Returns: Tuple of (full_key, full_value) with shape [batch, heads, text+image, head_dim] """ full_key = mx.concatenate([text_key, self.key_cache], axis=2) full_value = mx.concatenate([text_value, self.value_cache], axis=2) return full_key, full_value def reset(self) -> None: """Reset cache to zeros.""" self.key_cache = mx.zeros( (self.batch_size, self.num_heads, self.image_seq_len, self.head_dim), dtype=self._dtype, ) self.value_cache = mx.zeros( (self.batch_size, self.num_heads, self.image_seq_len, self.head_dim), dtype=self._dtype, ) ================================================ FILE: src/exo/worker/engines/image/pipeline/runner.py ================================================ from collections.abc import Iterator from dataclasses import dataclass from math import ceil from typing import Any, Optional, final import mlx.core as mx from mflux.models.common.config.config import Config from mflux.utils.exceptions import StopImageGenerationException from tqdm import tqdm from exo.shared.constants import EXO_TRACING_ENABLED from exo.shared.tracing import ( clear_trace_buffer, trace, ) from exo.shared.types.worker.shards import CfgShardMetadata, PipelineShardMetadata from exo.worker.engines.image.config import ImageModelConfig from exo.worker.engines.image.models.base import ( ModelAdapter, PromptData, RotaryEmbeddings, ) from exo.worker.engines.image.pipeline.block_wrapper import ( BlockWrapperMode, JointBlockWrapper, SingleBlockWrapper, ) @final @dataclass(frozen=True) class CfgBranch: positive: bool embeds: mx.array mask: mx.array | None pooled: mx.array | None cond_latents: mx.array | None def calculate_patch_heights( latent_height: int, num_patches: int ) -> tuple[list[int], int]: patch_height = ceil(latent_height / num_patches) actual_num_patches = ceil(latent_height / patch_height) patch_heights = [patch_height] * (actual_num_patches - 1) last_height = latent_height - patch_height * (actual_num_patches - 1) patch_heights.append(last_height) return patch_heights, actual_num_patches def calculate_token_indices( patch_heights: list[int], latent_width: int ) -> list[tuple[int, int]]: tokens_per_row = latent_width token_ranges: list[tuple[int, int]] = [] cumulative_height = 0 for h in patch_heights: start_token = tokens_per_row * cumulative_height end_token = tokens_per_row * (cumulative_height + h) token_ranges.append((start_token, end_token)) cumulative_height += h return token_ranges class DiffusionRunner: """Orchestrates the diffusion loop for image generation. In distributed mode, it implements PipeFusion with: - Sync pipeline for initial timesteps (full image, all devices in lockstep) - Async pipeline for later timesteps (patches processed independently) """ def __init__( self, config: ImageModelConfig, adapter: ModelAdapter[Any, Any], group: Optional[mx.distributed.Group], shard_metadata: PipelineShardMetadata | CfgShardMetadata, num_patches: Optional[int] = None, ): self.config = config self.adapter = adapter self.group = group self._init_cfg_topology(shard_metadata) self.num_patches = ( num_patches if num_patches else max(1, self.pipeline_world_size) ) self.total_joint = config.joint_block_count self.total_single = config.single_block_count self.total_layers = config.total_blocks self._guidance_override: float | None = None self._compute_assigned_blocks() def _init_cfg_topology( self, shard_metadata: PipelineShardMetadata | CfgShardMetadata ) -> None: """Initialize CFG and pipeline topology from shard metadata. Both CfgShardMetadata and PipelineShardMetadata represent pipeline parallel execution. CFG adds a second parallel pipeline for negative prompt processing, but within each pipeline group the communication pattern is identical. """ if self.group is None: # Single node - no distributed communication self.rank = 0 self.world_size = 1 self.start_layer = 0 self.end_layer = self.config.total_blocks self.cfg_rank = 0 self.cfg_world_size = 1 self.cfg_parallel = False self.pipeline_rank = 0 self.pipeline_world_size = 1 self.next_pipeline_rank: int | None = None self.prev_pipeline_rank: int | None = None self.cfg_peer_rank: int | None = None self.first_pipeline_rank: int = 0 self.last_pipeline_rank: int = 0 return # Common fields from base metadata self.rank = shard_metadata.device_rank self.world_size = shard_metadata.world_size self.start_layer = shard_metadata.start_layer self.end_layer = shard_metadata.end_layer if isinstance(shard_metadata, CfgShardMetadata): # CFG parallel: two independent pipelines self.cfg_rank = shard_metadata.cfg_rank self.cfg_world_size = shard_metadata.cfg_world_size self.cfg_parallel = True self.pipeline_rank = shard_metadata.pipeline_rank self.pipeline_world_size = shard_metadata.pipeline_world_size else: # Pure pipeline: single pipeline group, sequential CFG self.cfg_rank = 0 self.cfg_world_size = 1 self.cfg_parallel = False self.pipeline_rank = shard_metadata.device_rank self.pipeline_world_size = shard_metadata.world_size # Pipeline neighbor computation (same logic for both types) is_first = self.pipeline_rank == 0 is_last = self.pipeline_rank == self.pipeline_world_size - 1 self.next_pipeline_rank = ( None if is_last else self._device_rank_for(self.cfg_rank, self.pipeline_rank + 1) ) self.prev_pipeline_rank = ( None if is_first else self._device_rank_for(self.cfg_rank, self.pipeline_rank - 1) ) # CFG peer is the corresponding last stage in the other CFG group if self.cfg_parallel and is_last: other_cfg_rank = 1 - self.cfg_rank self.cfg_peer_rank = self._device_rank_for( other_cfg_rank, self.pipeline_rank ) else: self.cfg_peer_rank = None # First/last pipeline ranks for ring communication (latent broadcast) self.first_pipeline_rank = self._device_rank_for(self.cfg_rank, 0) self.last_pipeline_rank = self._device_rank_for( self.cfg_rank, self.pipeline_world_size - 1 ) def _device_rank_for(self, cfg_rank: int, pipeline_rank: int) -> int: """Convert (cfg_rank, pipeline_rank) to device_rank in the ring topology. Ring layout: [cfg0_pipe0, cfg0_pipe1, ..., cfg1_pipeN-1, cfg1_pipeN-2, ..., cfg1_pipe0] Group 0 is in ascending order, group 1 is reversed so last stages are neighbors. """ if not self.cfg_parallel: return pipeline_rank if cfg_rank == 0: return pipeline_rank else: return self.world_size - 1 - pipeline_rank def _compute_assigned_blocks(self) -> None: """Determine which joint/single blocks this stage owns.""" start = self.start_layer end = self.end_layer if end <= self.total_joint: self.joint_start = start self.joint_end = end self.single_start = 0 self.single_end = 0 elif start >= self.total_joint: self.joint_start = 0 self.joint_end = 0 self.single_start = start - self.total_joint self.single_end = end - self.total_joint else: self.joint_start = start self.joint_end = self.total_joint self.single_start = 0 self.single_end = end - self.total_joint self.has_joint_blocks = self.joint_end > self.joint_start self.has_single_blocks = self.single_end > self.single_start self.owns_concat_stage = self.has_joint_blocks and ( self.has_single_blocks or self.end_layer == self.total_joint ) # Wrappers created lazily on first forward (need text_seq_len) self.joint_block_wrappers: list[JointBlockWrapper[Any]] | None = None self.single_block_wrappers: list[SingleBlockWrapper[Any]] | None = None self._wrappers_initialized = False self._current_text_seq_len: int | None = None @property def is_first_stage(self) -> bool: return self.pipeline_rank == 0 @property def is_last_stage(self) -> bool: return self.pipeline_rank == self.pipeline_world_size - 1 @property def is_distributed(self) -> bool: return self.group is not None def _get_effective_guidance_scale(self) -> float | None: if self._guidance_override is not None: return self._guidance_override return self.config.guidance_scale def _get_cfg_branches(self, prompt_data: PromptData) -> Iterator[CfgBranch]: """Yield the CFG branches this node should process. - No CFG: yields one branch (positive) - CFG parallel: yields one branch (our assigned branch) - Sequential CFG: yields two branches (positive, then negative) """ if not self.adapter.needs_cfg: embeds, mask, pooled, cond = prompt_data.get_cfg_branch_data(positive=True) yield CfgBranch( positive=True, embeds=embeds, mask=mask, pooled=pooled, cond_latents=cond, ) elif self.cfg_parallel: positive = self.cfg_rank == 0 embeds, mask, pooled, cond = prompt_data.get_cfg_branch_data(positive) yield CfgBranch( positive=positive, embeds=embeds, mask=mask, pooled=pooled, cond_latents=cond, ) else: pos_embeds, pos_mask, pos_pooled, pos_cond = ( prompt_data.get_cfg_branch_data(positive=True) ) yield CfgBranch( positive=True, embeds=pos_embeds, mask=pos_mask, pooled=pos_pooled, cond_latents=pos_cond, ) neg_embeds, neg_mask, neg_pooled, neg_cond = ( prompt_data.get_cfg_branch_data(positive=False) ) yield CfgBranch( positive=False, embeds=neg_embeds, mask=neg_mask, pooled=neg_pooled, cond_latents=neg_cond, ) def _combine_cfg_results(self, results: list[tuple[bool, mx.array]]) -> mx.array: if len(results) == 1: positive, noise = results[0] if self.cfg_parallel and self.is_last_stage: # TODO(ciaran): try to remove mx.eval(noise) return self._exchange_and_apply_guidance(noise, positive) return noise noise_neg = next(n for p, n in results if not p) noise_pos = next(n for p, n in results if p) return self._apply_guidance(noise_pos, noise_neg) def _exchange_and_apply_guidance( self, noise: mx.array, is_positive: bool ) -> mx.array: assert self.group is not None assert self.cfg_peer_rank is not None if is_positive: noise = mx.distributed.send(noise, self.cfg_peer_rank, group=self.group) mx.async_eval(noise) noise_neg = mx.distributed.recv_like( noise, self.cfg_peer_rank, group=self.group ) mx.eval(noise_neg) noise_pos = noise else: noise_pos = mx.distributed.recv_like( noise, self.cfg_peer_rank, group=self.group ) mx.eval(noise_pos) noise = mx.distributed.send(noise, self.cfg_peer_rank, group=self.group) mx.async_eval(noise) noise_neg = noise return self._apply_guidance(noise_pos, noise_neg) def _apply_guidance(self, noise_pos: mx.array, noise_neg: mx.array) -> mx.array: scale = self._get_effective_guidance_scale() assert scale is not None return self.adapter.apply_guidance(noise_pos, noise_neg, scale) def _ensure_wrappers( self, text_seq_len: int, encoder_hidden_states_mask: mx.array | None = None, ) -> None: """Lazily create block wrappers on first forward pass. Wrappers need text_seq_len which is only known after prompt encoding. Re-initializes if text_seq_len changes (e.g., warmup vs real generation). """ if self._wrappers_initialized and self._current_text_seq_len == text_seq_len: return self.joint_block_wrappers = self.adapter.get_joint_block_wrappers( text_seq_len=text_seq_len, encoder_hidden_states_mask=encoder_hidden_states_mask, ) self.single_block_wrappers = self.adapter.get_single_block_wrappers( text_seq_len=text_seq_len, ) self._wrappers_initialized = True self._current_text_seq_len = text_seq_len def _reset_all_caches(self) -> None: """Reset KV caches on all wrappers for a new generation.""" if self.joint_block_wrappers: for wrapper in self.joint_block_wrappers: wrapper.reset_cache() if self.single_block_wrappers: for wrapper in self.single_block_wrappers: wrapper.reset_cache() def _set_text_seq_len(self, text_seq_len: int) -> None: if self.joint_block_wrappers: for wrapper in self.joint_block_wrappers: wrapper.set_text_seq_len(text_seq_len) if self.single_block_wrappers: for wrapper in self.single_block_wrappers: wrapper.set_text_seq_len(text_seq_len) def _calculate_capture_steps( self, partial_images: int, init_time_step: int, num_inference_steps: int, ) -> set[int]: """Calculate which timesteps should produce partial images. Places the first partial after step 1 for fast initial feedback, then evenly spaces remaining partials with equal gaps between them and from the last partial to the final image. Args: partial_images: Number of partial images to capture init_time_step: Starting timestep (for img2img this may not be 0) num_inference_steps: Total inference steps Returns: Set of timestep indices to capture """ if partial_images <= 0: return set() total_steps = num_inference_steps - init_time_step if total_steps <= 1: return set() if partial_images >= total_steps - 1: return set(range(init_time_step, num_inference_steps - 1)) capture_steps: set[int] = set() first_capture = init_time_step + 1 capture_steps.add(first_capture) if partial_images == 1: return capture_steps final_step = num_inference_steps - 1 remaining_range = final_step - first_capture for i in range(1, partial_images): step_idx = first_capture + int(i * remaining_range / partial_images) capture_steps.add(step_idx) return capture_steps def generate_image( self, runtime_config: Config, prompt: str, seed: int, partial_images: int = 0, guidance_override: float | None = None, negative_prompt: str | None = None, num_sync_steps: int = 1, ): """Primary entry point for image generation. Orchestrates the full generation flow: 1. Create runtime config 2. Create initial latents 3. Encode prompt 4. Run diffusion loop (yielding partials if requested) 5. Decode to image Args: settings: Generation config (steps, height, width) prompt: Text prompt seed: Random seed partial_images: Number of intermediate images to yield (0 for none) guidance_override: Optional override for guidance scale (CFG) Yields: Partial images as (GeneratedImage, partial_index, total_partials) tuples Final GeneratedImage """ self._guidance_override = guidance_override latents = self.adapter.create_latents(seed, runtime_config) prompt_data = self.adapter.encode_prompt(prompt, negative_prompt) capture_steps = self._calculate_capture_steps( partial_images=partial_images, init_time_step=runtime_config.init_time_step, num_inference_steps=runtime_config.num_inference_steps, ) diffusion_gen = self._run_diffusion_loop( latents=latents, prompt_data=prompt_data, runtime_config=runtime_config, seed=seed, prompt=prompt, capture_steps=capture_steps, num_sync_steps=num_sync_steps, ) partial_index = 0 total_partials = len(capture_steps) if capture_steps: try: while True: partial_latents, _step = next(diffusion_gen) if self.is_last_stage: partial_image = self.adapter.decode_latents( partial_latents, runtime_config, seed, prompt ) yield (partial_image, partial_index, total_partials) partial_index += 1 except StopIteration as e: latents = e.value # pyright: ignore[reportAny] else: try: while True: next(diffusion_gen) except StopIteration as e: latents = e.value # pyright: ignore[reportAny] if self.is_last_stage: yield self.adapter.decode_latents(latents, runtime_config, seed, prompt) # pyright: ignore[reportAny] def _run_diffusion_loop( self, latents: mx.array, prompt_data: PromptData, runtime_config: Config, seed: int, prompt: str, num_sync_steps: int, capture_steps: set[int] | None = None, ): if capture_steps is None: capture_steps = set() self._reset_all_caches() clear_trace_buffer() time_steps = tqdm(range(runtime_config.num_inference_steps)) ctx = self.adapter.model.callbacks.start( # pyright: ignore[reportAny] seed=seed, prompt=prompt, config=runtime_config ) ctx.before_loop( # pyright: ignore[reportAny] latents=latents, ) for t in time_steps: try: latents = self._diffusion_step( t=t, config=runtime_config, latents=latents, prompt_data=prompt_data, num_sync_steps=num_sync_steps, ) ctx.in_loop( # pyright: ignore[reportAny] t=t, latents=latents, time_steps=time_steps, ) mx.eval(latents) if t in capture_steps and self.is_last_stage: yield (latents, t) except KeyboardInterrupt: # noqa: PERF203 ctx.interruption(t=t, latents=latents) # pyright: ignore[reportAny] raise StopImageGenerationException( f"Stopping image generation at step {t + 1}/{len(time_steps)}" ) from None ctx.after_loop(latents=latents) # pyright: ignore[reportAny] return latents def _forward_pass( self, latents: mx.array, prompt_embeds: mx.array, pooled_prompt_embeds: mx.array, t: int, config: Config, encoder_hidden_states_mask: mx.array | None = None, cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]] | None = None, conditioning_latents: mx.array | None = None, kontext_image_ids: mx.array | None = None, ) -> mx.array: """Run a single forward pass through the transformer. Args: latents: Input latents (already scaled by caller) prompt_embeds: Text embeddings pooled_prompt_embeds: Pooled text embeddings (Flux) or placeholder (Qwen) t: Current timestep config: Runtime configuration encoder_hidden_states_mask: Attention mask for text (Qwen) cond_image_grid: Conditioning image grid dimensions (Qwen edit) conditioning_latents: Conditioning latents for edit mode kontext_image_ids: Position IDs for Kontext conditioning (Flux Kontext) Returns: Noise prediction tensor """ text_seq_len = prompt_embeds.shape[1] self._ensure_wrappers(text_seq_len, encoder_hidden_states_mask) if self.joint_block_wrappers and encoder_hidden_states_mask is not None: for wrapper in self.joint_block_wrappers: wrapper.set_encoder_mask(encoder_hidden_states_mask) scaled_latents = config.scheduler.scale_model_input(latents, t) # pyright: ignore[reportAny] # For edit mode: concatenate with conditioning latents original_latent_tokens: int = scaled_latents.shape[1] # pyright: ignore[reportAny] if conditioning_latents is not None: scaled_latents = mx.concatenate( [scaled_latents, conditioning_latents], axis=1 ) hidden_states, encoder_hidden_states = self.adapter.compute_embeddings( scaled_latents, prompt_embeds ) text_embeddings = self.adapter.compute_text_embeddings( t, config, pooled_prompt_embeds, hidden_states=hidden_states ) rotary_embeddings = self.adapter.compute_rotary_embeddings( prompt_embeds, config, encoder_hidden_states_mask=encoder_hidden_states_mask, cond_image_grid=cond_image_grid, kontext_image_ids=kontext_image_ids, ) assert self.joint_block_wrappers is not None for wrapper in self.joint_block_wrappers: encoder_hidden_states, hidden_states = wrapper( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, text_embeddings=text_embeddings, rotary_embeddings=rotary_embeddings, ) if self.joint_block_wrappers: hidden_states = self.adapter.merge_streams( hidden_states, encoder_hidden_states ) assert self.single_block_wrappers is not None for wrapper in self.single_block_wrappers: hidden_states = wrapper( hidden_states=hidden_states, text_embeddings=text_embeddings, rotary_embeddings=rotary_embeddings, ) # Extract image portion hidden_states = hidden_states[:, text_seq_len:, ...] # For edit mode: extract only the generated portion (exclude conditioning latents) if conditioning_latents is not None: hidden_states = hidden_states[:, :original_latent_tokens, ...] return self.adapter.final_projection(hidden_states, text_embeddings) def _diffusion_step( self, t: int, config: Config, latents: mx.array, prompt_data: PromptData, num_sync_steps: int, ) -> mx.array: if self.group is None: return self._single_node_step(t, config, latents, prompt_data) elif ( self.pipeline_world_size == 1 or t < config.init_time_step + num_sync_steps ): with trace(name=f"sync {t}", rank=self.rank, category="sync"): return self._sync_pipeline_step( t, config, latents, prompt_data, ) else: with trace(name=f"async {t}", rank=self.rank, category="async"): return self._async_pipeline_step( t, config, latents, prompt_data, is_first_async_step=t == config.init_time_step + num_sync_steps, ) def _single_node_step( self, t: int, config: Config, latents: mx.array, prompt_data: PromptData, ) -> mx.array: cond_image_grid = prompt_data.cond_image_grid kontext_image_ids = prompt_data.kontext_image_ids results: list[tuple[bool, mx.array]] = [] for branch in self._get_cfg_branches(prompt_data): # Reset caches before each branch to ensure no state contamination self._reset_all_caches() pooled_embeds = ( branch.pooled if branch.pooled is not None else branch.embeds ) noise = self._forward_pass( latents, branch.embeds, pooled_embeds, t=t, config=config, encoder_hidden_states_mask=branch.mask, cond_image_grid=cond_image_grid, conditioning_latents=branch.cond_latents, kontext_image_ids=kontext_image_ids, ) results.append((branch.positive, noise)) noise = self._combine_cfg_results(results) return config.scheduler.step(noise=noise, timestep=t, latents=latents) # pyright: ignore[reportAny] def _create_patches( self, latents: mx.array, config: Config, ) -> tuple[list[mx.array], list[tuple[int, int]]]: latent_height = config.height // 16 latent_width = config.width // 16 patch_heights, _ = calculate_patch_heights(latent_height, self.num_patches) token_indices = calculate_token_indices(patch_heights, latent_width) patch_latents = [latents[:, start:end, :] for start, end in token_indices] return patch_latents, token_indices def _run_sync_pass( self, t: int, config: Config, scaled_hidden_states: mx.array, prompt_embeds: mx.array, pooled_prompt_embeds: mx.array, encoder_hidden_states_mask: mx.array | None, cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]] | None, kontext_image_ids: mx.array | None, num_img_tokens: int, original_latent_tokens: int, conditioning_latents: mx.array | None, ) -> mx.array | None: hidden_states = scaled_hidden_states batch_size = hidden_states.shape[0] text_seq_len = prompt_embeds.shape[1] hidden_dim = self.adapter.hidden_dim dtype = scaled_hidden_states.dtype self._set_text_seq_len(text_seq_len) if self.joint_block_wrappers: for wrapper in self.joint_block_wrappers: wrapper.set_encoder_mask(encoder_hidden_states_mask) encoder_hidden_states: mx.array | None = None if self.is_first_stage: hidden_states, encoder_hidden_states = self.adapter.compute_embeddings( hidden_states, prompt_embeds ) text_embeddings = self.adapter.compute_text_embeddings( t, config, pooled_prompt_embeds, hidden_states=hidden_states ) image_rotary_embeddings = self.adapter.compute_rotary_embeddings( prompt_embeds, config, encoder_hidden_states_mask=encoder_hidden_states_mask, cond_image_grid=cond_image_grid, kontext_image_ids=kontext_image_ids, ) if self.has_joint_blocks: if not self.is_first_stage: assert self.prev_pipeline_rank is not None with trace( name=f"recv {self.prev_pipeline_rank}", rank=self.rank, category="comms", ): hidden_states = mx.distributed.recv( (batch_size, num_img_tokens, hidden_dim), dtype, self.prev_pipeline_rank, group=self.group, ) encoder_hidden_states = mx.distributed.recv( (batch_size, text_seq_len, hidden_dim), dtype, self.prev_pipeline_rank, group=self.group, ) mx.eval(hidden_states, encoder_hidden_states) assert self.joint_block_wrappers is not None assert encoder_hidden_states is not None with trace( name="joint_blocks", rank=self.rank, category="compute", ): for wrapper in self.joint_block_wrappers: wrapper.set_patch(BlockWrapperMode.CACHING) encoder_hidden_states, hidden_states = wrapper( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, text_embeddings=text_embeddings, rotary_embeddings=image_rotary_embeddings, ) if EXO_TRACING_ENABLED: mx.eval(encoder_hidden_states, hidden_states) if self.owns_concat_stage: assert encoder_hidden_states is not None concatenated = self.adapter.merge_streams( hidden_states, encoder_hidden_states ) if self.has_single_blocks or self.is_last_stage: hidden_states = concatenated else: assert self.next_pipeline_rank is not None with trace( name=f"send {self.next_pipeline_rank}", rank=self.rank, category="comms", ): concatenated = mx.distributed.send( concatenated, self.next_pipeline_rank, group=self.group ) mx.async_eval(concatenated) elif self.has_joint_blocks and not self.is_last_stage: assert encoder_hidden_states is not None assert self.next_pipeline_rank is not None with trace( name=f"send {self.next_pipeline_rank}", rank=self.rank, category="comms", ): hidden_states = mx.distributed.send( hidden_states, self.next_pipeline_rank, group=self.group ) encoder_hidden_states = mx.distributed.send( encoder_hidden_states, self.next_pipeline_rank, group=self.group ) mx.async_eval(hidden_states, encoder_hidden_states) if self.has_single_blocks: if not self.owns_concat_stage and not self.is_first_stage: assert self.prev_pipeline_rank is not None with trace( name=f"recv {self.prev_pipeline_rank}", rank=self.rank, category="comms", ): hidden_states = mx.distributed.recv( (batch_size, text_seq_len + num_img_tokens, hidden_dim), dtype, self.prev_pipeline_rank, group=self.group, ) mx.eval(hidden_states) assert self.single_block_wrappers is not None with trace( name="single blocks", rank=self.rank, category="compute", ): for wrapper in self.single_block_wrappers: wrapper.set_patch(BlockWrapperMode.CACHING) hidden_states = wrapper( hidden_states=hidden_states, text_embeddings=text_embeddings, rotary_embeddings=image_rotary_embeddings, ) if EXO_TRACING_ENABLED: mx.eval(hidden_states) if not self.is_last_stage: assert self.next_pipeline_rank is not None with trace( name=f"send {self.next_pipeline_rank}", rank=self.rank, category="comms", ): hidden_states = mx.distributed.send( hidden_states, self.next_pipeline_rank, group=self.group ) mx.async_eval(hidden_states) hidden_states = hidden_states[:, text_seq_len:, ...] if conditioning_latents is not None: hidden_states = hidden_states[:, :original_latent_tokens, ...] if self.is_last_stage: return self.adapter.final_projection(hidden_states, text_embeddings) return None def _sync_pipeline_step( self, t: int, config: Config, hidden_states: mx.array, prompt_data: PromptData, ) -> mx.array: prev_latents = hidden_states cond_image_grid = prompt_data.cond_image_grid kontext_image_ids = prompt_data.kontext_image_ids scaled_hidden_states = config.scheduler.scale_model_input(hidden_states, t) # pyright: ignore[reportAny] original_latent_tokens: int = scaled_hidden_states.shape[1] # pyright: ignore[reportAny] results: list[tuple[bool, mx.array]] = [] for branch in self._get_cfg_branches(prompt_data): pooled_embeds = ( branch.pooled if branch.pooled is not None else branch.embeds ) cond_latents = branch.cond_latents if cond_latents is not None: num_img_tokens: int = original_latent_tokens + cond_latents.shape[1] else: num_img_tokens = original_latent_tokens step_latents: mx.array = scaled_hidden_states # pyright: ignore[reportAny] if self.is_first_stage and cond_latents is not None: step_latents = mx.concatenate([step_latents, cond_latents], axis=1) text_seq_len = branch.embeds.shape[1] self._ensure_wrappers(text_seq_len, branch.mask) noise = self._run_sync_pass( t, config, step_latents, branch.embeds, pooled_embeds, branch.mask, cond_image_grid, kontext_image_ids, num_img_tokens, original_latent_tokens, cond_latents, ) if self.is_last_stage: assert noise is not None results.append((branch.positive, noise)) if self.is_last_stage: noise = self._combine_cfg_results(results) hidden_states = config.scheduler.step( # pyright: ignore[reportAny] noise=noise, timestep=t, latents=prev_latents ) if not self.is_first_stage: hidden_states = mx.distributed.send( hidden_states, self.first_pipeline_rank, group=self.group ) mx.async_eval(hidden_states) elif self.is_first_stage: hidden_states = mx.distributed.recv_like( prev_latents, src=self.last_pipeline_rank, group=self.group ) mx.eval(hidden_states) else: hidden_states = prev_latents return hidden_states def _async_pipeline_step( self, t: int, config: Config, latents: mx.array, prompt_data: PromptData, is_first_async_step: bool, ) -> mx.array: patch_latents, token_indices = self._create_patches(latents, config) cond_image_grid = prompt_data.cond_image_grid kontext_image_ids = prompt_data.kontext_image_ids prev_patch_latents = [p for p in patch_latents] encoder_hidden_states: mx.array | None = None for patch_idx in range(len(patch_latents)): patch = patch_latents[patch_idx] if ( self.is_first_stage and not self.is_last_stage and not is_first_async_step ): with trace( name=f"recv {self.last_pipeline_rank}", rank=self.rank, category="comms", ): patch = mx.distributed.recv_like( patch, src=self.last_pipeline_rank, group=self.group ) mx.eval(patch) results: list[tuple[bool, mx.array]] = [] for branch in self._get_cfg_branches(prompt_data): pooled_embeds = ( branch.pooled if branch.pooled is not None else branch.embeds ) text_seq_len = branch.embeds.shape[1] self._ensure_wrappers(text_seq_len, branch.mask) self._set_text_seq_len(text_seq_len) if self.joint_block_wrappers: for wrapper in self.joint_block_wrappers: wrapper.set_encoder_mask(branch.mask) text_embeddings = self.adapter.compute_text_embeddings( t, config, pooled_embeds ) image_rotary_embeddings = self.adapter.compute_rotary_embeddings( branch.embeds, config, encoder_hidden_states_mask=branch.mask, cond_image_grid=cond_image_grid, kontext_image_ids=kontext_image_ids, ) noise, encoder_hidden_states = self._run_single_patch_pass( patch=patch, patch_idx=patch_idx, token_indices=token_indices[patch_idx], prompt_embeds=branch.embeds, text_embeddings=text_embeddings, image_rotary_embeddings=image_rotary_embeddings, encoder_hidden_states=encoder_hidden_states, ) if self.is_last_stage: assert noise is not None results.append((branch.positive, noise)) if self.is_last_stage: noise = self._combine_cfg_results(results) patch_latents[patch_idx] = config.scheduler.step( # pyright: ignore[reportAny] noise=noise, timestep=t, latents=prev_patch_latents[patch_idx], ) if not self.is_first_stage and t != config.num_inference_steps - 1: with trace( name=f"send {self.first_pipeline_rank}", rank=self.rank, category="comms", ): patch_latents[patch_idx] = mx.distributed.send( patch_latents[patch_idx], self.first_pipeline_rank, group=self.group, ) mx.async_eval(patch_latents[patch_idx]) return mx.concatenate(patch_latents, axis=1) def _run_single_patch_pass( self, patch: mx.array, patch_idx: int, token_indices: tuple[int, int], prompt_embeds: mx.array, text_embeddings: mx.array, image_rotary_embeddings: RotaryEmbeddings, encoder_hidden_states: mx.array | None, ) -> tuple[mx.array | None, mx.array | None]: """Process a single patch through the forward pipeline. Handles stage-to-stage communication (stage i -> stage i+1). Ring communication (last stage -> first stage) is handled by the caller. Args: patch: The patch latents to process patch_idx: Index of this patch (0-indexed) token_indices: (start_token, end_token) for this patch prompt_embeds: Text embeddings (for compute_embeddings on first stage) text_embeddings: Precomputed text embeddings image_rotary_embeddings: Precomputed rotary embeddings encoder_hidden_states: Encoder hidden states (passed between patches) Returns: (noise_prediction, encoder_hidden_states) - noise is None for non-last stages """ start_token, end_token = token_indices batch_size = patch.shape[0] text_seq_len = prompt_embeds.shape[1] hidden_dim = self.adapter.hidden_dim if self.has_joint_blocks: if not self.is_first_stage: assert self.prev_pipeline_rank is not None patch_len = patch.shape[1] with trace( name=f"recv {self.prev_pipeline_rank}", rank=self.rank, category="comms", ): patch = mx.distributed.recv( (batch_size, patch_len, hidden_dim), patch.dtype, self.prev_pipeline_rank, group=self.group, ) mx.eval(patch) if patch_idx == 0: with trace( name=f"recv {self.prev_pipeline_rank}", rank=self.rank, category="comms", ): encoder_hidden_states = mx.distributed.recv( (batch_size, text_seq_len, hidden_dim), patch.dtype, self.prev_pipeline_rank, group=self.group, ) mx.eval(encoder_hidden_states) if self.is_first_stage: patch, encoder_hidden_states = self.adapter.compute_embeddings( patch, prompt_embeds ) assert self.joint_block_wrappers is not None assert encoder_hidden_states is not None with trace( name=f"joint patch {patch_idx}", rank=self.rank, category="compute", ): for wrapper in self.joint_block_wrappers: wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token) encoder_hidden_states, patch = wrapper( hidden_states=patch, encoder_hidden_states=encoder_hidden_states, text_embeddings=text_embeddings, rotary_embeddings=image_rotary_embeddings, ) if EXO_TRACING_ENABLED: mx.eval(encoder_hidden_states, patch) if self.owns_concat_stage: assert encoder_hidden_states is not None patch_concat = self.adapter.merge_streams(patch, encoder_hidden_states) if self.has_single_blocks or self.is_last_stage: patch = patch_concat else: assert self.next_pipeline_rank is not None with trace( name=f"send {self.next_pipeline_rank}", rank=self.rank, category="comms", ): patch_concat = mx.distributed.send( patch_concat, self.next_pipeline_rank, group=self.group ) mx.async_eval(patch_concat) elif self.has_joint_blocks and not self.is_last_stage: assert self.next_pipeline_rank is not None with trace( name=f"send {self.next_pipeline_rank}", rank=self.rank, category="comms", ): patch = mx.distributed.send( patch, self.next_pipeline_rank, group=self.group ) mx.async_eval(patch) if patch_idx == 0: assert encoder_hidden_states is not None with trace( name=f"send {self.next_pipeline_rank}", rank=self.rank, category="comms", ): encoder_hidden_states = mx.distributed.send( encoder_hidden_states, self.next_pipeline_rank, group=self.group ) mx.async_eval(encoder_hidden_states) if self.has_single_blocks: if not self.owns_concat_stage and not self.is_first_stage: assert self.prev_pipeline_rank is not None patch_len = patch.shape[1] with trace( name=f"recv {self.prev_pipeline_rank}", rank=self.rank, category="comms", ): patch = mx.distributed.recv( (batch_size, text_seq_len + patch_len, hidden_dim), patch.dtype, self.prev_pipeline_rank, group=self.group, ) mx.eval(patch) assert self.single_block_wrappers is not None with trace( name=f"single patch {patch_idx}", rank=self.rank, category="compute", ): for wrapper in self.single_block_wrappers: wrapper.set_patch(BlockWrapperMode.PATCHED, start_token, end_token) patch = wrapper( hidden_states=patch, text_embeddings=text_embeddings, rotary_embeddings=image_rotary_embeddings, ) if EXO_TRACING_ENABLED: mx.eval(patch) if not self.is_last_stage: assert self.next_pipeline_rank is not None with trace( name=f"send {self.next_pipeline_rank}", rank=self.rank, category="comms", ): patch = mx.distributed.send( patch, self.next_pipeline_rank, group=self.group ) mx.async_eval(patch) noise: mx.array | None = None if self.is_last_stage: patch_img_only = patch[:, text_seq_len:, :] noise = self.adapter.final_projection(patch_img_only, text_embeddings) return noise, encoder_hidden_states ================================================ FILE: src/exo/worker/engines/mlx/__init__.py ================================================ ================================================ FILE: src/exo/worker/engines/mlx/auto_parallel.py ================================================ import os import threading from abc import ABC, abstractmethod from collections.abc import Callable from functools import partial from inspect import signature from typing import TYPE_CHECKING, Any, Literal, Protocol, cast import mlx.core as mx import mlx.nn as nn from mlx.nn.layers.distributed import ( shard_inplace, shard_linear, sum_gradients, ) from mlx_lm.models.base import ( scaled_dot_product_attention, # pyright: ignore[reportUnknownVariableType] ) from mlx_lm.models.cache import ArraysCache, KVCache from mlx_lm.models.deepseek_v3 import DeepseekV3MLP from mlx_lm.models.deepseek_v3 import Model as DeepseekV3Model from mlx_lm.models.deepseek_v32 import DeepseekV32MLP from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model from mlx_lm.models.glm4_moe import Model as Glm4MoeModel from mlx_lm.models.glm4_moe import MoE from mlx_lm.models.glm4_moe_lite import Glm4MoeLiteDecoderLayer, Glm4MoeLiteMLP from mlx_lm.models.glm4_moe_lite import Model as GLM4MoeLiteModel from mlx_lm.models.gpt_oss import GptOssMoeModel from mlx_lm.models.gpt_oss import Model as GptOssModel from mlx_lm.models.kimi_k25 import Model as KimiK25Model from mlx_lm.models.llama import Model as LlamaModel from mlx_lm.models.minimax import MiniMaxAttention from mlx_lm.models.minimax import Model as MiniMaxModel from mlx_lm.models.ministral3 import Model as Ministral3Model from mlx_lm.models.nemotron_h import Model as NemotronHModel from mlx_lm.models.nemotron_h import ( NemotronHAttention, NemotronHMamba2Mixer, NemotronHMoE, ) from mlx_lm.models.nemotron_h import NemotronHModel as NemotronHInnerModel from mlx_lm.models.qwen3_5 import DecoderLayer as Qwen3_5DecoderLayer from mlx_lm.models.qwen3_5 import Model as Qwen3_5TextModel from mlx_lm.models.qwen3_5 import Qwen3_5TextModel as Qwen3_5TextModelInner from mlx_lm.models.qwen3_5 import SparseMoeBlock as Qwen3_5SparseMoeBlock from mlx_lm.models.qwen3_5_moe import Model as Qwen3_5MoeModel from mlx_lm.models.qwen3_moe import Model as Qwen3MoeModel from mlx_lm.models.qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeSparseMoeBlock from mlx_lm.models.qwen3_next import Model as Qwen3NextModel from mlx_lm.models.qwen3_next import ( Qwen3NextDecoderLayer, Qwen3NextGatedDeltaNet, Qwen3NextSparseMoeBlock, ) from mlx_lm.models.qwen3_next import Qwen3NextModel as Qwen3NextInnerModel from mlx_lm.models.step3p5 import Model as Step35Model from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel from exo.shared.logging import logger from exo.shared.types.worker.shards import PipelineShardMetadata if TYPE_CHECKING: from mlx_lm.models.cache import Cache TimeoutCallback = Callable[[], None] LayerLoadedCallback = Callable[[int, int], None] # (layers_loaded, total_layers) _pending_prefill_sends: list[tuple[mx.array, int, mx.distributed.Group]] = [] def flush_prefill_sends() -> None: for output, dst, group in _pending_prefill_sends: sent = mx.distributed.send(output, dst, group=group) mx.async_eval(sent) _pending_prefill_sends.clear() def clear_prefill_sends() -> None: # Discard pending sends (e.g. on cancellation). _pending_prefill_sends.clear() def eval_with_timeout( mlx_item: Any, # pyright: ignore[reportAny] timeout_seconds: float = 60.0, on_timeout: TimeoutCallback | None = None, ) -> None: """Evaluate MLX item with a hard timeout. If on_timeout callback is provided, it will be called before terminating the process. This allows the runner to send a failure event before exit. """ completed = threading.Event() def watchdog() -> None: if not completed.wait(timeout=timeout_seconds): logger.error( f"mlx_item evaluation timed out after {timeout_seconds:.0f}s. " "This may indicate an issue with FAST_SYNCH and tensor parallel sharding. " "Terminating process." ) if on_timeout is not None: on_timeout() os._exit(1) watchdog_thread = threading.Thread(target=watchdog, daemon=True) watchdog_thread.start() try: mx.eval(mlx_item) # pyright: ignore[reportAny] finally: completed.set() class _LayerCallable(Protocol): """Structural type that any compatible layer must satisfy. We require a single positional input of type ``mx.array`` and an ``mx.array`` output, while permitting arbitrary *args / **kwargs so this protocol matches the vast majority of `mlx.nn.Module` subclasses. """ def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array: ... class CustomMlxLayer(nn.Module): """Base class for replacing an MLX layer with a custom implementation.""" def __init__(self, original_layer: _LayerCallable): super().__init__() dict.__setitem__(self, "_original_layer", original_layer) # pyright: ignore[reportUnknownMemberType] @property def original_layer(self) -> _LayerCallable: return cast(_LayerCallable, self["_original_layer"]) # Calls __getattr__ for any attributes not found on nn.Module (e.g. use_sliding) if not TYPE_CHECKING: def __getattr__(self, name): try: return super().__getattr__(name) except AttributeError: original_layer = cast(_LayerCallable, self["_original_layer"]) return getattr(original_layer, name) class PipelineFirstLayer(CustomMlxLayer): def __init__( self, original_layer: _LayerCallable, r: int, group: mx.distributed.Group, ): super().__init__(original_layer) self.r: int = r self.group = group self.is_prefill: bool = False def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array: if self.r != 0: # We want to avoid GPU timeout errors by evalling the distributed operation # so that it stays on CPU, which does not have a timeout. mx.eval(x) x = mx.distributed.recv_like(x, (self.r - 1), group=self.group) mx.eval(x) return self.original_layer(x, *args, **kwargs) class PipelineLastLayer(CustomMlxLayer): def __init__( self, original_layer: _LayerCallable, r: int, s: int, group: mx.distributed.Group, ): super().__init__(original_layer) self.r: int = r self.s: int = s self.group = group self.original_layer_signature = signature(self.original_layer.__call__) self.is_prefill: bool = False self.queue_sends: bool = False def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array: cache = self.original_layer_signature.bind_partial( x, *args, **kwargs ).arguments.get("cache", None) output: mx.array = self.original_layer(x, *args, **kwargs) # Eval layer output to materialize it before send — this splits the graph # so the send is isolated and the receiving rank's recv can complete. mx.eval(output) if self.r != self.s - 1: if self.queue_sends: _pending_prefill_sends.append( (output, (self.r + 1) % self.s, self.group) ) else: output = mx.distributed.send( output, (self.r + 1) % self.s, group=self.group ) if cache is not None: # CacheList (used by MLA models like DeepSeekV32, GLM MoE DSA) # doesn't have .keys directly; access via first sub-cache. _cache = cache[0] if hasattr(cache, "caches") else cache # type: ignore if hasattr(_cache, "keys"): # pyright: ignore[reportAny] _cache.keys = mx.depends(_cache.keys, output) # type: ignore mx.eval(output) if cache is not None and hasattr(_cache, "keys"): # type: ignore mx.eval(_cache.keys) # type: ignore if not self.is_prefill: output = mx.distributed.all_gather(output, group=self.group)[ -output.shape[0] : ] mx.eval(output) return output def set_pipeline_prefill(model: nn.Module, is_prefill: bool) -> None: for layer in model.layers: # type: ignore if isinstance(layer, (PipelineFirstLayer, PipelineLastLayer)): layer.is_prefill = is_prefill def set_pipeline_queue_sends(model: nn.Module, queue_sends: bool) -> None: for layer in model.layers: # type: ignore if isinstance(layer, PipelineLastLayer): layer.queue_sends = queue_sends def get_inner_model(model: nn.Module) -> nn.Module: inner = getattr(model, "model", None) if isinstance(inner, nn.Module): return inner inner = getattr(model, "transformer", None) if isinstance(inner, nn.Module): return inner inner = getattr(model, "language_model", None) if isinstance(inner, nn.Module): inner_inner = getattr(inner, "model", None) if isinstance(inner_inner, nn.Module): return inner_inner inner = getattr(model, "backbone", None) if isinstance(inner, nn.Module): return inner raise ValueError( "Model must either have a 'model', 'transformer', or 'backbone' attribute" ) def get_layers(inner_model_instance: nn.Module) -> list[_LayerCallable]: # Handle both model.layers and model.h cases layers: list[_LayerCallable] if hasattr(inner_model_instance, "layers"): layers = cast(list[_LayerCallable], inner_model_instance.layers) elif hasattr(inner_model_instance, "h"): layers = cast(list[_LayerCallable], inner_model_instance.h) else: raise ValueError("Model must have either a 'layers' or 'h' attribute") return layers def _patch_hybrid_cache( model: Qwen3_5TextModel | Qwen3NextModel | NemotronHModel, fa_idx: int, has_full_attn: bool, ssm_idx: int, has_linear: bool, ) -> None: # Hacks to make make_mask happy. original = model.make_cache def patched() -> list[ArraysCache | KVCache]: cache = original() if not has_full_attn: entry = cache[fa_idx] orig_make_mask = entry.make_mask entry.make_mask = lambda n, **_kw: orig_make_mask(n) # type: ignore if not has_linear: orig_ssm_make_mask = cache[ssm_idx].make_mask def _ssm_mask( n: int, **kw: bool | int | None ) -> mx.array | Literal["causal"] | None: return orig_ssm_make_mask(n, **kw) if kw else None cache[ssm_idx].make_mask = _ssm_mask # type: ignore return cache model.make_cache = patched def pipeline_auto_parallel( model: nn.Module, group: mx.distributed.Group, model_shard_meta: PipelineShardMetadata, on_layer_loaded: LayerLoadedCallback | None, ) -> nn.Module: """ Automatically parallelize a model across multiple devices. Args: model: The model to parallelize (must have a 'layers' or 'h' property) model_shard_meta: The metadata for the model shard Returns: The parallelized model """ inner_model_instance: nn.Module = get_inner_model(model) layers = get_layers(inner_model_instance) start_layer, end_layer = model_shard_meta.start_layer, model_shard_meta.end_layer device_rank, world_size = model_shard_meta.device_rank, model_shard_meta.world_size layers = layers[start_layer:end_layer] total = len(layers) for i, layer in enumerate(layers): mx.eval(layer) # type: ignore if on_layer_loaded is not None: on_layer_loaded(i, total) layers[0] = PipelineFirstLayer(layers[0], device_rank, group=group) layers[-1] = PipelineLastLayer( layers[-1], device_rank, world_size, group=group, ) if isinstance(inner_model_instance, GptOssMoeModel): inner_model_instance.layer_types = inner_model_instance.layer_types[ # type: ignore start_layer:end_layer ] # We can assume the model has at least one layer thanks to placement. # If a layer type doesn't exist, we can set it to 0. inner_model_instance.swa_idx = ( 0 if "sliding_attention" not in inner_model_instance.layer_types # type: ignore else inner_model_instance.layer_types.index( # type: ignore "sliding_attention" ) ) inner_model_instance.ga_idx = ( 0 if "full_attention" not in inner_model_instance.layer_types # type: ignore else inner_model_instance.layer_types.index( # type: ignore "full_attention" ) ) if isinstance(inner_model_instance, Step35InnerModel): inner_model_instance.num_layers = len(layers) sliding_layers = [ i for i, layer in enumerate(layers) if getattr(layer, "is_sliding", False) ] full_layers = [ i for i, layer in enumerate(layers) if not getattr(layer, "is_sliding", True) ] inner_model_instance._swa_idx = 0 if not sliding_layers else sliding_layers[0] inner_model_instance._full_idx = 0 if not full_layers else full_layers[0] if isinstance(inner_model_instance, (Qwen3_5TextModelInner, Qwen3NextInnerModel)): full_attn_layers = [ i for i, layer in enumerate(layers) if not getattr(layer, "is_linear", True) ] linear_layers = [ i for i, layer in enumerate(layers) if getattr(layer, "is_linear", False) ] inner_model_instance.fa_idx = full_attn_layers[0] if full_attn_layers else 0 inner_model_instance.ssm_idx = linear_layers[0] if linear_layers else 0 if not full_attn_layers or not linear_layers: _patch_hybrid_cache( cast(Qwen3_5TextModel | Qwen3NextModel, model), fa_idx=inner_model_instance.fa_idx, has_full_attn=bool(full_attn_layers), ssm_idx=inner_model_instance.ssm_idx, has_linear=bool(linear_layers), ) if isinstance(inner_model_instance, NemotronHInnerModel): # NemotronH uses block_type: "M" (Mamba/SSM), "*" (Attention), "E" (MoE), "-" (MLP) # Only "M" and "*" blocks have cache entries. # Recompute fa_idx and ssm_idx as cache-array indices for the shard's layers. cache_idx = 0 fa_idx: int | None = None ssm_idx: int | None = None for layer in layers: block_type = getattr(layer, "block_type", None) if block_type == "*": if fa_idx is None: fa_idx = cache_idx cache_idx += 1 elif block_type == "M": if ssm_idx is None: ssm_idx = cache_idx cache_idx += 1 has_attn = fa_idx is not None has_mamba = ssm_idx is not None inner_model_instance.fa_idx = fa_idx if fa_idx is not None else 0 inner_model_instance.ssm_idx = ssm_idx if ssm_idx is not None else 0 if not has_attn or not has_mamba: _patch_hybrid_cache( cast(NemotronHModel, model), fa_idx=inner_model_instance.fa_idx, has_full_attn=has_attn, ssm_idx=inner_model_instance.ssm_idx, has_linear=has_mamba, ) _set_layers(model, layers) assert isinstance(layers, list), ( "Expected a list of layers after auto-parallel initialisation" ) return patch_pipeline_model(model, group) def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T: # Patch __call__ on the model's class cls = model.__class__ original_call = cls.__call__ # type :ignore call_signature = signature(original_call) # type :ignore def patched_call( self: T, *args: object, **kwargs: object, ) -> mx.array: logits: mx.array = original_call(self, *args, **kwargs) # type: ignore cache = call_signature.bind_partial(self, *args, **kwargs).arguments.get( "cache", None ) # Add dependency to last cache entry to ensure distributed ops are evaluated if cache is not None: last = cache[-1] # type: ignore dep_cache = last[0] if hasattr(last, "caches") else last # type: ignore if hasattr(dep_cache, "keys") and dep_cache.keys is not None: # type: ignore dep_cache.keys = mx.depends(dep_cache.keys, logits) # type: ignore return logits cls.__call__ = patched_call return model def patch_tensor_model[T](model: T) -> T: """Patch model's __call__ to ensure distributed ops sync during inference.""" cls = model.__class__ original_call = cls.__call__ call_signature = signature(original_call) def patched_call( self: T, *args: object, **kwargs: object, ) -> mx.array: logits: mx.array = original_call(self, *args, **kwargs) # pyright: ignore[reportAny] cache = call_signature.bind_partial(self, *args, **kwargs).arguments.get( "cache", None ) # Add dependency to last cache entry to ensure distributed ops are evaluated if cache is not None and len(cache) > 0: # pyright: ignore[reportAny] last = cache[-1] # pyright: ignore[reportAny] dep_cache = last[0] if hasattr(last, "caches") else last # pyright: ignore[reportAny] if hasattr(dep_cache, "keys"): # type: ignore dep_cache.keys = mx.depends(dep_cache.keys, logits) # pyright: ignore[reportAny,reportUnknownMemberType] return logits cls.__call__ = patched_call return model def tensor_auto_parallel( model: nn.Module, group: mx.distributed.Group, timeout_seconds: float, on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, ) -> nn.Module: all_to_sharded_linear = partial( shard_linear, sharding="all-to-sharded", group=group, ) sharded_to_all_linear = partial( shard_linear, sharding="sharded-to-all", group=group, ) segments: int = 1 def _all_to_sharded(path: str, weight: mx.array): if path.endswith("bias"): logger.info(f"Sharding bias for {path} - all to sharded") return weight.ndim - 1, segments return max(weight.ndim - 2, 0), segments all_to_sharded_linear_in_place = partial( shard_inplace, sharding=_all_to_sharded, # type: ignore group=group, ) n = group.size() def _sharded_to_all(path: str, weight: mx.array): if path.endswith("bias"): logger.info(f"Sharding bias for {path} - sharded to all") weight /= n return None return -1, segments sharded_to_all_linear_in_place = partial( shard_inplace, sharding=_sharded_to_all, # type: ignore group=group, ) if isinstance(model, (LlamaModel, Ministral3Model)): tensor_parallel_sharding_strategy = LlamaShardingStrategy( group, all_to_sharded_linear, sharded_to_all_linear, all_to_sharded_linear_in_place, sharded_to_all_linear_in_place, ) elif isinstance(model, (DeepseekV3Model, DeepseekV32Model, KimiK25Model)): tensor_parallel_sharding_strategy = DeepSeekShardingStrategy( group, all_to_sharded_linear, sharded_to_all_linear, all_to_sharded_linear_in_place, sharded_to_all_linear_in_place, ) elif isinstance(model, MiniMaxModel): tensor_parallel_sharding_strategy = MiniMaxShardingStrategy( group, all_to_sharded_linear, sharded_to_all_linear, all_to_sharded_linear_in_place, sharded_to_all_linear_in_place, ) elif isinstance(model, GLM4MoeLiteModel): tensor_parallel_sharding_strategy = GLM4MoeLiteShardingStrategy( group, all_to_sharded_linear, sharded_to_all_linear, all_to_sharded_linear_in_place, sharded_to_all_linear_in_place, ) elif isinstance(model, Glm4MoeModel): tensor_parallel_sharding_strategy = Glm4MoeShardingStrategy( group, all_to_sharded_linear, sharded_to_all_linear, all_to_sharded_linear_in_place, sharded_to_all_linear_in_place, ) elif isinstance( model, (Qwen3MoeModel, Qwen3NextModel, Qwen3_5TextModel, Qwen3_5MoeModel) ): tensor_parallel_sharding_strategy = QwenShardingStrategy( group, all_to_sharded_linear, sharded_to_all_linear, all_to_sharded_linear_in_place, sharded_to_all_linear_in_place, ) elif isinstance(model, GptOssModel): tensor_parallel_sharding_strategy = GptOssShardingStrategy( group, all_to_sharded_linear, sharded_to_all_linear, all_to_sharded_linear_in_place, sharded_to_all_linear_in_place, ) elif isinstance(model, Step35Model): tensor_parallel_sharding_strategy = Step35ShardingStrategy( group, all_to_sharded_linear, sharded_to_all_linear, all_to_sharded_linear_in_place, sharded_to_all_linear_in_place, ) elif isinstance(model, NemotronHModel): tensor_parallel_sharding_strategy = NemotronHShardingStrategy( group, all_to_sharded_linear, sharded_to_all_linear, all_to_sharded_linear_in_place, sharded_to_all_linear_in_place, ) else: raise ValueError(f"Unsupported model type: {type(model)}") model = tensor_parallel_sharding_strategy.shard_model( model, timeout_seconds, on_timeout, on_layer_loaded ) return patch_tensor_model(model) class TensorParallelShardingStrategy(ABC): def __init__( self, group: mx.distributed.Group, all_to_sharded_linear: Callable[..., nn.Linear], sharded_to_all_linear: Callable[..., nn.Linear], all_to_sharded_linear_in_place: Callable[..., None], sharded_to_all_linear_in_place: Callable[..., None], ): self.all_to_sharded_linear = all_to_sharded_linear self.sharded_to_all_linear = sharded_to_all_linear self.all_to_sharded_linear_in_place = all_to_sharded_linear_in_place self.sharded_to_all_linear_in_place = sharded_to_all_linear_in_place self.group = group self.N = group.size() @abstractmethod def shard_model( self, model: nn.Module, timeout_seconds: float, on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, ) -> nn.Module: ... class LlamaShardingStrategy(TensorParallelShardingStrategy): def shard_model( self, model: nn.Module, timeout_seconds: float, on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, ) -> nn.Module: model = cast(LlamaModel, model) total = len(model.layers) for i, layer in enumerate(model.layers): # Force load weights before sharding to avoid FAST_SYNCH deadlock eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout) layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj) layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj) layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj) layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj) layer.self_attn.n_heads //= self.N if layer.self_attn.n_kv_heads is not None: layer.self_attn.n_kv_heads //= self.N layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj) layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj) layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj) mx.eval(layer) if on_layer_loaded is not None: on_layer_loaded(i, total) return model def _set_layers(model: nn.Module, layers: list[_LayerCallable]) -> None: inner_model_instance = get_inner_model(model) if hasattr(inner_model_instance, "layers"): inner_model_instance.layers = layers # Update DeepSeek V3 specific parameters when layers are shrunk if isinstance( model, (DeepseekV3Model, DeepseekV32Model, Glm4MoeModel, KimiK25Model) ) and hasattr(inner_model_instance, "num_layers"): logger.info( f"Setting num_layers to {len(layers)} for model {model.model.__class__.__name__}" ) inner_model_instance.start_idx = 0 inner_model_instance.end_idx = len(layers) inner_model_instance.num_layers = len(layers) elif isinstance(model, Qwen3MoeModel): logger.info( f"Setting num_hidden_layers to {len(layers)} for model {model.model.__class__.__name__}" ) inner_model_instance.num_hidden_layers = len(layers) elif hasattr(inner_model_instance, "h"): inner_model_instance.h = layers else: raise ValueError("Model must have either a 'layers' or 'h' attribute") class DeepSeekShardingStrategy(TensorParallelShardingStrategy): def shard_model( self, model: nn.Module, timeout_seconds: float, on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, ) -> nn.Module: model = cast(DeepseekV3Model, model) total = len(model.layers) for i, layer in enumerate(model.layers): eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout) # Shard the self attention if layer.self_attn.q_lora_rank is None: layer.self_attn.q_proj = self.all_to_sharded_linear( layer.self_attn.q_proj ) else: layer.self_attn.q_b_proj = self.all_to_sharded_linear( layer.self_attn.q_b_proj ) layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj) layer.self_attn.num_heads //= self.N # Logic from upstream mlx num_heads = layer.self_attn.num_heads sh = self.group.rank() * num_heads eh = sh + num_heads def shard_heads(w: mx.array, sh: int = sh, eh: int = eh) -> mx.array: return w[sh:eh] layer.self_attn.embed_q.apply(shard_heads) layer.self_attn.unembed_out.apply(shard_heads) # Shard the MLP if isinstance(layer.mlp, (DeepseekV3MLP, DeepseekV32MLP)): layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj) layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj) layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj) # Shard the MoE. else: if getattr(layer.mlp, "shared_experts", None) is not None: self.all_to_sharded_linear_in_place( layer.mlp.shared_experts.gate_proj ) self.sharded_to_all_linear_in_place( layer.mlp.shared_experts.down_proj ) self.all_to_sharded_linear_in_place( layer.mlp.shared_experts.up_proj ) self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj) self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj) self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj) layer.mlp = ShardedMoE(layer.mlp) # type: ignore layer.mlp.sharding_group = self.group mx.eval(layer) if on_layer_loaded is not None: on_layer_loaded(i, total) return model class ShardedMoE(CustomMlxLayer): """Wraps any MoE layer with distributed sum_gradients / all_sum.""" def __init__(self, layer: _LayerCallable): super().__init__(layer) self.sharding_group: mx.distributed.Group | None = None def __call__(self, x: mx.array) -> mx.array: if self.sharding_group is not None: x = sum_gradients(self.sharding_group)(x) y = self.original_layer.__call__(x) if self.sharding_group is not None: y = mx.distributed.all_sum(y, group=self.sharding_group) return y class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy): def shard_model( self, model: nn.Module, timeout_seconds: float, on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, ) -> nn.Module: model = cast(GLM4MoeLiteModel, model) total = len(model.layers) # type: ignore for i, layer in enumerate(model.layers): # type: ignore layer = cast(Glm4MoeLiteDecoderLayer, layer) eval_with_timeout( layer.parameters(), timeout_seconds / total, on_timeout, ) if layer.self_attn.q_lora_rank is None: # type: ignore layer.self_attn.q_proj = self.all_to_sharded_linear( layer.self_attn.q_proj ) else: layer.self_attn.q_b_proj = self.all_to_sharded_linear( layer.self_attn.q_b_proj ) layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj) layer.self_attn.num_heads //= self.N # Logic from upstream mlx num_heads = layer.self_attn.num_heads sh = self.group.rank() * num_heads eh = sh + num_heads def shard_heads(w: mx.array, sh: int = sh, eh: int = eh) -> mx.array: return w[sh:eh] layer.self_attn.embed_q.apply(shard_heads) layer.self_attn.unembed_out.apply(shard_heads) if isinstance(layer.mlp, Glm4MoeLiteMLP): layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj) layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj) layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj) else: if getattr(layer.mlp, "shared_experts", None) is not None: self.all_to_sharded_linear_in_place( layer.mlp.shared_experts.gate_proj ) self.sharded_to_all_linear_in_place( layer.mlp.shared_experts.down_proj ) self.all_to_sharded_linear_in_place( layer.mlp.shared_experts.up_proj ) self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj) self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj) self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj) layer.mlp = ShardedMoE(layer.mlp) # type: ignore layer.mlp.sharding_group = self.group # type: ignore mx.eval(layer) if on_layer_loaded is not None: on_layer_loaded(i, total) return model class WrappedMiniMaxAttention(CustomMlxLayer): def __init__(self, layer: _LayerCallable, group: mx.distributed.Group): super().__init__(layer) self.group = group def __call__( self, x: mx.array, mask: mx.array | None = None, cache: "Cache | None" = None, ) -> mx.array: batch_dim, seq_dim, _ = x.shape self._original_layer = cast(MiniMaxAttention, self.original_layer) # type: ignore queries: mx.array = self._original_layer.q_proj(x) keys: mx.array = self._original_layer.k_proj(x) values: mx.array = self._original_layer.v_proj(x) if getattr(self, "use_qk_norm", False): q_dim = queries.shape[-1] k_dim = keys.shape[-1] n = self.group.size() qk = mx.concatenate( [queries, keys], axis=-1 ) # (batch_dim, seq_dim, q_dim + k_dim) qk = mx.distributed.all_gather( qk, group=self.group ) # (n*batch_dim, seq_dim, q_dim + k_dim) qk = qk.reshape(n, batch_dim, seq_dim, q_dim + k_dim).transpose(1, 2, 0, 3) queries = qk[..., :q_dim].reshape( batch_dim, seq_dim, -1 ) # (batch_dim, seq_dim, n * q_dim) keys = qk[..., q_dim:].reshape( batch_dim, seq_dim, -1 ) # (batch_dim, seq_dim, n * k_dim) queries = self._original_layer.q_norm(queries) keys = self._original_layer.k_norm(keys) # Split back and take this rank's portion queries = mx.split(queries, n, axis=-1)[self.group.rank()] keys = mx.split(keys, n, axis=-1)[self.group.rank()] queries = queries.reshape( batch_dim, seq_dim, self._original_layer.num_attention_heads, -1 ).transpose(0, 2, 1, 3) keys = keys.reshape( batch_dim, seq_dim, self._original_layer.num_key_value_heads, -1 ).transpose(0, 2, 1, 3) values = values.reshape( batch_dim, seq_dim, self._original_layer.num_key_value_heads, -1 ).transpose(0, 2, 1, 3) if cache is not None: queries = self._original_layer.rope(queries, offset=cache.offset) keys = self._original_layer.rope(keys, offset=cache.offset) keys, values = cache.update_and_fetch(keys, values) else: queries = self._original_layer.rope(queries) keys = self._original_layer.rope(keys) output = scaled_dot_product_attention( queries, keys, values, cache=cache, scale=self._original_layer.scale, # type: ignore mask=mask, ) output = output.transpose(0, 2, 1, 3).reshape(batch_dim, seq_dim, -1) return self._original_layer.o_proj(output) class MiniMaxShardingStrategy(TensorParallelShardingStrategy): def shard_model( self, model: nn.Module, timeout_seconds: float, on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, ) -> nn.Module: model = cast(MiniMaxModel, model) total = len(model.layers) for i, layer in enumerate(model.layers): eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout) # Shard the self attention layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj) layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj) layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj) layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj) layer.self_attn.num_attention_heads //= self.N layer.self_attn.num_key_value_heads //= self.N layer.self_attn = WrappedMiniMaxAttention(layer.self_attn, self.group) # pyright: ignore[reportAttributeAccessIssue,reportArgumentType] # Shard the MoE. self.all_to_sharded_linear_in_place( layer.block_sparse_moe.switch_mlp.gate_proj ) self.sharded_to_all_linear_in_place( layer.block_sparse_moe.switch_mlp.down_proj ) self.all_to_sharded_linear_in_place( layer.block_sparse_moe.switch_mlp.up_proj ) layer.block_sparse_moe = ShardedMoE(layer.block_sparse_moe) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] layer.block_sparse_moe.sharding_group = self.group # pyright: ignore[reportAttributeAccessIssue] mx.eval(layer) if on_layer_loaded is not None: on_layer_loaded(i, total) return model class QwenShardingStrategy(TensorParallelShardingStrategy): def shard_model( self, model: nn.Module, timeout_seconds: float, on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, ) -> nn.Module: model = cast( Qwen3MoeModel | Qwen3NextModel | Qwen3_5TextModel | Qwen3_5MoeModel, model ) total = len(model.layers) for i, layer in enumerate(model.layers): eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout) # Shard the self attention if isinstance(layer, Qwen3MoeDecoderLayer): layer.self_attn.q_proj = self.all_to_sharded_linear( layer.self_attn.q_proj ) layer.self_attn.k_proj = self.all_to_sharded_linear( layer.self_attn.k_proj ) layer.self_attn.v_proj = self.all_to_sharded_linear( layer.self_attn.v_proj ) layer.self_attn.o_proj = self.sharded_to_all_linear( layer.self_attn.o_proj ) layer.self_attn.n_heads //= self.N layer.self_attn.n_kv_heads //= self.N else: assert isinstance(layer, (Qwen3NextDecoderLayer, Qwen3_5DecoderLayer)) if hasattr(layer, "linear_attn"): linear_attn = layer.linear_attn if isinstance(linear_attn, Qwen3NextGatedDeltaNet): # Qwen3-Next: combined projections linear_attn.in_proj_qkvz = self.all_to_sharded_linear( linear_attn.in_proj_qkvz ) linear_attn.in_proj_ba = self.all_to_sharded_linear( linear_attn.in_proj_ba ) else: # Qwen3.5: separate projections # in_proj_qkv has sections [q(key_dim), k(key_dim), v(value_dim)] # that must be split section-aware, not as a contiguous block key_dim = linear_attn.key_dim value_dim = linear_attn.value_dim linear_attn.in_proj_qkv = shard_linear( linear_attn.in_proj_qkv, "all-to-sharded", segments=[key_dim, key_dim + key_dim], group=self.group, ) linear_attn.in_proj_z = self.all_to_sharded_linear( linear_attn.in_proj_z ) linear_attn.in_proj_b = self.all_to_sharded_linear( linear_attn.in_proj_b ) linear_attn.in_proj_a = self.all_to_sharded_linear( linear_attn.in_proj_a ) linear_attn.out_proj = self.sharded_to_all_linear( linear_attn.out_proj ) # Shard conv1d: depthwise conv with non-contiguous channel slicing. # Channel layout is [q(key_dim), k(key_dim), v(value_dim)]. # Each rank takes its head-slice from each of the three sections. rank = self.group.rank() key_dim = linear_attn.key_dim value_dim = linear_attn.value_dim key_dim_shard = key_dim // self.N value_dim_shard = value_dim // self.N q_idx = mx.arange(rank * key_dim_shard, (rank + 1) * key_dim_shard) k_idx = mx.arange( key_dim + rank * key_dim_shard, key_dim + (rank + 1) * key_dim_shard, ) v_idx = mx.arange( 2 * key_dim + rank * value_dim_shard, 2 * key_dim + (rank + 1) * value_dim_shard, ) conv_indices = mx.concatenate([q_idx, k_idx, v_idx]) linear_attn.conv1d.weight = linear_attn.conv1d.weight[conv_indices] new_conv_dim = key_dim_shard * 2 + value_dim_shard linear_attn.conv1d.groups = new_conv_dim num_v_shard = linear_attn.num_v_heads // self.N v_start = rank * num_v_shard v_end = v_start + num_v_shard linear_attn.A_log = linear_attn.A_log[v_start:v_end] linear_attn.dt_bias = linear_attn.dt_bias[v_start:v_end] linear_attn.num_k_heads //= self.N linear_attn.num_v_heads //= self.N linear_attn.key_dim = ( linear_attn.head_k_dim * linear_attn.num_k_heads ) linear_attn.value_dim = ( linear_attn.head_v_dim * linear_attn.num_v_heads ) linear_attn.conv_dim = ( linear_attn.key_dim * 2 + linear_attn.value_dim ) else: layer.self_attn.q_proj = self.all_to_sharded_linear( layer.self_attn.q_proj ) layer.self_attn.k_proj = self.all_to_sharded_linear( layer.self_attn.k_proj ) layer.self_attn.v_proj = self.all_to_sharded_linear( layer.self_attn.v_proj ) layer.self_attn.o_proj = self.sharded_to_all_linear( layer.self_attn.o_proj ) layer.self_attn.num_attention_heads //= self.N layer.self_attn.num_key_value_heads //= self.N # Shard the MoE. if isinstance( layer.mlp, ( Qwen3MoeSparseMoeBlock, Qwen3NextSparseMoeBlock, Qwen3_5SparseMoeBlock, ), ): self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj) self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj) self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj) if isinstance( layer.mlp, (Qwen3NextSparseMoeBlock, Qwen3_5SparseMoeBlock) ): self.all_to_sharded_linear_in_place( layer.mlp.shared_expert.gate_proj ) self.sharded_to_all_linear_in_place( layer.mlp.shared_expert.down_proj ) self.all_to_sharded_linear_in_place(layer.mlp.shared_expert.up_proj) layer.mlp = ShardedMoE(layer.mlp) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] layer.mlp.sharding_group = self.group # Shard the MLP else: layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj) layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj) layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj) mx.eval(layer) if on_layer_loaded is not None: on_layer_loaded(i, total) return model class Glm4MoeShardingStrategy(TensorParallelShardingStrategy): def shard_model( self, model: nn.Module, timeout_seconds: float, on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, ) -> nn.Module: model = cast(Glm4MoeModel, model) total = len(model.layers) for i, layer in enumerate(model.layers): eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout) layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj) layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj) layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj) layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj) layer.self_attn.n_heads //= self.N layer.self_attn.n_kv_heads //= self.N if isinstance(layer.mlp, MoE): self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj) self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj) self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj) if getattr(layer.mlp, "shared_experts", None) is not None: self.all_to_sharded_linear_in_place( layer.mlp.shared_experts.gate_proj ) self.sharded_to_all_linear_in_place( layer.mlp.shared_experts.down_proj ) self.all_to_sharded_linear_in_place( layer.mlp.shared_experts.up_proj ) layer.mlp = ShardedMoE(layer.mlp) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] layer.mlp.sharding_group = self.group else: layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj) layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj) layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj) mx.eval(layer) if on_layer_loaded is not None: on_layer_loaded(i, total) return model class GptOssShardingStrategy(TensorParallelShardingStrategy): def shard_model( self, model: nn.Module, timeout_seconds: float, on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, ) -> nn.Module: model = cast(GptOssMoeModel, model) total = len(model.layers) for i, layer in enumerate(model.layers): eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout) layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj) layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj) layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj) layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj) layer.self_attn.num_attention_heads //= self.N layer.self_attn.num_key_value_heads //= self.N layer.self_attn.num_key_value_groups = ( layer.self_attn.num_attention_heads // layer.self_attn.num_key_value_heads ) layer.self_attn.sinks = layer.self_attn.sinks[ layer.self_attn.num_attention_heads * self.group.rank() : layer.self_attn.num_attention_heads * (self.group.rank() + 1) ] self.all_to_sharded_linear_in_place(layer.mlp.experts.gate_proj) self.sharded_to_all_linear_in_place(layer.mlp.experts.down_proj) self.all_to_sharded_linear_in_place(layer.mlp.experts.up_proj) layer.mlp = ShardedMoE(layer.mlp) # type: ignore layer.mlp.sharding_group = self.group # pyright: ignore[reportAttributeAccessIssue] mx.eval(layer) if on_layer_loaded is not None: on_layer_loaded(i, total) return model class Step35ShardingStrategy(TensorParallelShardingStrategy): def shard_model( self, model: nn.Module, timeout_seconds: float, on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, ) -> nn.Module: model = cast(Step35Model, model) total = len(model.layers) for i, layer in enumerate(model.layers): eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout) layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj) layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj) layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj) layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj) layer.self_attn.num_heads //= self.N layer.self_attn.num_kv_heads //= self.N if getattr(layer.self_attn, "use_head_wise_attn_gate", False): layer.self_attn.g_proj = self.all_to_sharded_linear( layer.self_attn.g_proj ) if isinstance(layer.mlp, Step35MLP): layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj) layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj) layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj) else: layer.mlp.sharding_group = self.group self.all_to_sharded_linear_in_place(layer.mlp.share_expert.gate_proj) self.all_to_sharded_linear_in_place(layer.mlp.share_expert.up_proj) self.sharded_to_all_linear_in_place(layer.mlp.share_expert.down_proj) self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj) self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj) self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj) mx.eval(layer) if on_layer_loaded is not None: on_layer_loaded(i, total) return model class NemotronHShardingStrategy(TensorParallelShardingStrategy): def shard_model( self, model: nn.Module, timeout_seconds: float, on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, ) -> nn.Module: model = cast(NemotronHModel, model) rank = self.group.rank() total = len(model.layers) for i, layer in enumerate(model.layers): eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout) mixer = layer.mixer if isinstance(mixer, NemotronHAttention): mixer.q_proj = self.all_to_sharded_linear(mixer.q_proj) mixer.k_proj = self.all_to_sharded_linear(mixer.k_proj) mixer.v_proj = self.all_to_sharded_linear(mixer.v_proj) mixer.o_proj = self.sharded_to_all_linear(mixer.o_proj) mixer.num_heads //= self.N mixer.num_key_value_heads //= self.N elif isinstance(mixer, NemotronHMamba2Mixer): self._shard_mamba2_mixer(mixer, rank) elif isinstance(mixer, NemotronHMoE): # Shard routed experts (SwitchMLP uses fc1/fc2) self.all_to_sharded_linear_in_place(mixer.switch_mlp.fc1) self.sharded_to_all_linear_in_place(mixer.switch_mlp.fc2) # Shard shared expert in-place (no all-reduce — ShardedMoE handles that) if hasattr(mixer, "shared_experts"): self.all_to_sharded_linear_in_place(mixer.shared_experts.up_proj) self.sharded_to_all_linear_in_place(mixer.shared_experts.down_proj) mixer = ShardedMoE(mixer) # pyright: ignore[reportArgumentType] mixer.sharding_group = self.group layer.mixer = mixer # pyright: ignore[reportAttributeAccessIssue] mx.eval(layer) if on_layer_loaded is not None: on_layer_loaded(i, total) return model def _shard_mamba2_mixer(self, mixer: NemotronHMamba2Mixer, rank: int) -> None: """Shard the Mamba2 mixer along the head dimension.""" world_size = self.N num_heads = mixer.num_heads head_dim = mixer.head_dim n_groups = mixer.n_groups ssm_state_size = mixer.ssm_state_size intermediate_size = mixer.intermediate_size # = num_heads * head_dim # Per-rank sizes heads_per_rank = num_heads // world_size groups_per_rank = n_groups // world_size is_per_rank = heads_per_rank * head_dim bc_per_rank = groups_per_rank * ssm_state_size # === in_proj: output layout is [gate:IS | conv_ssm:IS | B:NG*SS | C:NG*SS | dt:NH] === gate_start = 0 conv_ssm_start = intermediate_size b_start = 2 * intermediate_size c_start = b_start + n_groups * ssm_state_size dt_start = c_start + n_groups * ssm_state_size # Build index tensor for this rank's slice of each section gate_idx = mx.arange( gate_start + rank * is_per_rank, gate_start + (rank + 1) * is_per_rank ) conv_ssm_idx = mx.arange( conv_ssm_start + rank * is_per_rank, conv_ssm_start + (rank + 1) * is_per_rank, ) b_idx = mx.arange( b_start + rank * bc_per_rank, b_start + (rank + 1) * bc_per_rank ) c_idx = mx.arange( c_start + rank * bc_per_rank, c_start + (rank + 1) * bc_per_rank ) dt_idx = mx.arange( dt_start + rank * heads_per_rank, dt_start + (rank + 1) * heads_per_rank ) indices = mx.concatenate([gate_idx, conv_ssm_idx, b_idx, c_idx, dt_idx]) mixer.in_proj.weight = mixer.in_proj.weight[indices] # === out_proj: input is intermediate_size (sharded) → hidden_size (reduce) === mixer.out_proj = self.sharded_to_all_linear(mixer.out_proj) # === conv1d: depthwise conv on conv_dim channels === # conv_dim layout: [ssm_hidden:IS | B:NG*SS | C:NG*SS] conv_ssm_idx_local = mx.arange(rank * is_per_rank, (rank + 1) * is_per_rank) conv_b_idx = mx.arange( intermediate_size + rank * bc_per_rank, intermediate_size + (rank + 1) * bc_per_rank, ) conv_c_idx = mx.arange( intermediate_size + n_groups * ssm_state_size + rank * bc_per_rank, intermediate_size + n_groups * ssm_state_size + (rank + 1) * bc_per_rank, ) conv_indices = mx.concatenate([conv_ssm_idx_local, conv_b_idx, conv_c_idx]) mixer.conv1d.weight = mixer.conv1d.weight[conv_indices] new_conv_dim = is_per_rank + 2 * bc_per_rank mixer.conv1d.groups = new_conv_dim if mixer.conv1d.bias is not None: mixer.conv1d.bias = mixer.conv1d.bias[conv_indices] # === Per-head parameters === h_start = rank * heads_per_rank h_end = h_start + heads_per_rank mixer.dt_bias = mixer.dt_bias[h_start:h_end] mixer.A_log = mixer.A_log[h_start:h_end] mixer.D = mixer.D[h_start:h_end] # === Norm: weight is intermediate_size === mixer.norm.weight = mixer.norm.weight[ rank * is_per_rank : (rank + 1) * is_per_rank ] # === Update dimensions === mixer.num_heads = heads_per_rank mixer.n_groups = groups_per_rank mixer.intermediate_size = is_per_rank mixer.conv_dim = new_conv_dim mixer.heads_per_group = heads_per_rank // groups_per_rank ================================================ FILE: src/exo/worker/engines/mlx/cache.py ================================================ import os from copy import deepcopy import mlx.core as mx import psutil from mlx_lm.models.cache import ( ArraysCache, CacheList, KVCache, QuantizedKVCache, RotatingKVCache, ) from mlx_lm.tokenizer_utils import TokenizerWrapper from exo.shared.types.memory import Memory from exo.shared.types.mlx import KVCacheType, Model from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS from exo.worker.runner.bootstrap import logger # Fraction of device memory above which LRU eviction kicks in. # Smaller machines need more aggressive eviction. def _default_memory_threshold() -> float: total_gb = Memory.from_bytes(psutil.virtual_memory().total).in_gb if total_gb >= 128: return 0.85 if total_gb >= 64: return 0.80 if total_gb >= 32: return 0.75 return 0.70 _MEMORY_THRESHOLD = float( os.environ.get("EXO_MEMORY_THRESHOLD", _default_memory_threshold()) ) class CacheSnapshot: """Snapshot of states at a known token position.""" def __init__( self, states: list[RotatingKVCache | ArraysCache | None], token_count: int ): self.states = states self.token_count = token_count def snapshot_ssm_states(cache: KVCacheType) -> CacheSnapshot: states: list[ArraysCache | RotatingKVCache | None] = [] for c in cache: if isinstance(c, (ArraysCache, RotatingKVCache)): states.append(deepcopy(c)) else: states.append(None) token_count = cache_length(cache) return CacheSnapshot(states=states, token_count=token_count) def _find_nearest_snapshot( snapshots: list[CacheSnapshot], target_token_count: int, ) -> CacheSnapshot | None: best: CacheSnapshot | None = None for snap in snapshots: if snap.token_count <= target_token_count and ( best is None or snap.token_count > best.token_count ): best = snap return best def has_non_kv_caches(cache: KVCacheType) -> bool: """Check if a cache contains any ArraysCache (SSM) entries.""" return any(isinstance(c, (ArraysCache, RotatingKVCache)) for c in cache) class KVPrefixCache: def __init__(self, group: mx.distributed.Group | None): self.prompts: list[mx.array] = [] # mx array of tokens (ints) self.caches: list[KVCacheType] = [] self._snapshots: list[list[CacheSnapshot] | None] = [] self._last_used: list[int] = [] # monotonic counter of last access per entry self._access_counter: int = 0 self._group = group def clear(self): """Clear all cached prompts and caches.""" self.prompts.clear() self.caches.clear() self._snapshots.clear() self._last_used.clear() def add_kv_cache( self, prompt_tokens: mx.array, cache: KVCacheType, ssm_snapshots: list[CacheSnapshot] | None = None, ): """Add a new cache entry. Evicts LRU entries if memory is high.""" self._evict_if_needed() self.prompts.append(prompt_tokens) self.caches.append(deepcopy(cache)) self._snapshots.append(ssm_snapshots) self._access_counter += 1 self._last_used.append(self._access_counter) logger.info(f"KV cache added: {len(prompt_tokens)} tokens") def update_kv_cache( self, index: int, prompt_tokens: mx.array, cache: KVCacheType, snapshots: list[CacheSnapshot] | None, restore_pos: int, ): """Update an existing cache entry in-place.""" old_snapshots = self._snapshots[index] merged: list[CacheSnapshot] = [] if old_snapshots: merged = [s for s in old_snapshots if s.token_count <= restore_pos] if snapshots: merged.extend(snapshots) self.prompts[index] = prompt_tokens self.caches[index] = deepcopy(cache) self._snapshots[index] = merged or None self._access_counter += 1 self._last_used[index] = self._access_counter logger.info(f"KV cache updated (index {index}): {len(prompt_tokens)} tokens") def _get_snapshot( self, entry_index: int, target_token_count: int ) -> tuple[int, CacheSnapshot | None]: if not has_non_kv_caches(self.caches[entry_index]): return target_token_count, None snapshots = self._snapshots[entry_index] if not snapshots: return 0, None snap = _find_nearest_snapshot(snapshots, target_token_count) if snap is not None: return snap.token_count, snap return 0, None def get_kv_cache( self, model: Model, prompt_tokens: mx.array, ) -> tuple[KVCacheType, mx.array, int | None]: """Get KV cache for prompt, returning remaining tokens to prefill. Returns: Tuple of (cache, remaining_tokens, matched_index) where: - cache: KV cache to use for generation - remaining_tokens: tokens that still need prefilling - matched_index: index of the matched entry (None if no match) For models with SSM layers (which are ArraysCache in mlx), the cache is trimmed to the nearest SSM snapshot position at or before the match point for correctness. Same for rotating KV Cache. """ max_length = len(prompt_tokens) best_index: int | None = None best_length = 0 is_exact = False # Find best cache match for i, cached_prompt in enumerate(self.prompts): length = get_prefix_length(prompt_tokens, cached_prompt) if length >= max_length - 1: best_index, best_length = i, length is_exact = True break if length > best_length: best_index, best_length = i, length if best_index is None: return make_kv_cache(model), prompt_tokens, None # For exact match: trim to max_length-1 so remaining has the last token # For partial match: trim to best_length, remaining has suffix to prefill # This ensures stream_generate always has at least one token to start with has_ssm = has_non_kv_caches(self.caches[best_index]) target = (max_length - 1) if is_exact and not has_ssm else best_length restore_pos, restore_snap = self._get_snapshot(best_index, target) # No usable snapshot — need fresh cache if restore_snap is None and has_ssm: return make_kv_cache(model), prompt_tokens, None prompt_cache = deepcopy(self.caches[best_index]) cached_length = cache_length(self.caches[best_index]) tokens_to_trim = cached_length - restore_pos if tokens_to_trim > 0: trim_cache(prompt_cache, tokens_to_trim, restore_snap) # Reset cache offset to match trimmed length for c in prompt_cache: if hasattr(c, "offset"): c.offset = restore_pos self._access_counter += 1 self._last_used[best_index] = self._access_counter remaining = prompt_tokens[restore_pos:] return prompt_cache, remaining, best_index def _evict_if_needed(self): """Evict least recently used entries while memory usage is high.""" if len(self.caches) == 0: return # Evict LRU entries until below threshold while ( len(self.caches) > 0 and self.get_memory_used_percentage() > _MEMORY_THRESHOLD ): lru_index = self._last_used.index(min(self._last_used)) evicted_tokens = len(self.prompts[lru_index]) self.prompts.pop(lru_index) self.caches.pop(lru_index) self._snapshots.pop(lru_index) self._last_used.pop(lru_index) logger.info( f"KV cache evicted LRU entry ({evicted_tokens} tokens) due to memory usage" ) def get_memory_used_percentage(self) -> float: local_pressure: float = get_memory_used_percentage() if self._group is None: return local_pressure all_pressure = mx.distributed.all_gather( mx.array([local_pressure], dtype=mx.float32), group=self._group, ) # .item() evals. max_pressure = float(mx.max(all_pressure).item()) return max_pressure def trim_cache( cache: KVCacheType, num_tokens: int, snapshot: CacheSnapshot | None = None, ) -> None: for i, c in enumerate(cache): if isinstance(c, (ArraysCache, RotatingKVCache)): if snapshot is not None and snapshot.states[i] is not None: cache[i] = deepcopy(snapshot.states[i]) # type: ignore else: c.state = [None] * len(c.state) else: c.trim(num_tokens) def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array: """Encode a prompt string to token array. For chat-templated prompts (which have their own structure markers like <|im_user|>, <|im_middle|>, etc.), we should NOT add BOS/EOS tokens as that would corrupt the prompt structure. """ # Chat templates define their own structure - don't add BOS/EOS prompt_tokens = tokenizer.encode(prompt, add_special_tokens=False) return mx.array(prompt_tokens) def _entry_length( c: KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList, ) -> int: # Use .offset attribute which KVCache types have (len() not implemented in older QuantizedKVCache). if hasattr(c, "offset"): return c.offset # For CacheList if hasattr(c, "size"): return int(c.size()) # type: ignore return 0 def cache_length(cache: KVCacheType) -> int: """Get the number of tokens in a KV cache.""" return max(_entry_length(c) for c in cache) def get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int: """Find the length of the common prefix between two token arrays.""" n = min(int(prompt.shape[0]), int(cached_prompt.shape[0])) if n == 0: return 0 equal = mx.equal(prompt[:n], cached_prompt[:n]).astype(mx.int32) prefix_mask = mx.cumprod(equal) # stays 1 until first mismatch, then 0 forever return int(mx.sum(prefix_mask).item()) def get_available_memory() -> Memory: mem: int = psutil.virtual_memory().available return Memory.from_bytes(mem) def get_memory_used_percentage() -> float: mem = psutil.virtual_memory() # percent is 0-100 return float(mem.percent / 100) def make_kv_cache( model: Model, max_kv_size: int | None = None, keep: int = 0 ) -> KVCacheType: assert hasattr(model, "layers") if hasattr(model, "make_cache"): logger.info("Using MLX LM's make cache") return model.make_cache() # type: ignore if max_kv_size is None: if KV_CACHE_BITS is None: logger.info("Using default KV cache") return [KVCache() for _ in model.layers] else: logger.info("Using quantized KV cache") return [ QuantizedKVCache(group_size=CACHE_GROUP_SIZE, bits=KV_CACHE_BITS) for _ in model.layers ] else: logger.info(f"Using rotating KV cache with {max_kv_size=} with {keep=}") return [RotatingKVCache(max_size=max_kv_size, keep=keep) for _ in model.layers] ================================================ FILE: src/exo/worker/engines/mlx/constants.py ================================================ # TODO: Do we want so many constants? # I think we want a lot of these as parameters? KV_GROUP_SIZE: int | None = 32 KV_BITS: int | None = None ATTENTION_KV_BITS: int | None = 4 MAX_TOKENS: int = 32168 MAX_KV_SIZE: int | None = 3200 KEEP_KV_SIZE: int | None = 1600 QUANTIZE_MODEL_MODE: str | None = "affine" CACHE_GROUP_SIZE: int = 64 KV_CACHE_BITS: int | None = None DEFAULT_TOP_LOGPROBS: int = 5 # TODO: We should really make this opt-in, but Kimi requires trust_remote_code=True TRUST_REMOTE_CODE: bool = True ================================================ FILE: src/exo/worker/engines/mlx/dsml_encoding.py ================================================ import json import re from typing import Any from mlx_lm.chat_templates import deepseek_v32 from exo.api.types import ToolCallItem BOS_TOKEN: str = deepseek_v32.bos_token EOS_TOKEN: str = deepseek_v32.eos_token DSML_TOKEN: str = deepseek_v32.dsml_token THINKING_START: str = deepseek_v32.thinking_start_token THINKING_END: str = deepseek_v32.thinking_end_token USER_TOKEN = "<\uff5cUser\uff5c>" ASSISTANT_TOKEN = "<\uff5cAssistant\uff5c>" TOOL_CALLS_START = f"<{DSML_TOKEN}function_calls>" TOOL_CALLS_END = f"" encode_messages = deepseek_v32.encode_messages _INVOKE_PATTERN = re.compile( rf"<{re.escape(DSML_TOKEN)}invoke\s+name=\"([^\"]+)\">" rf"(.*?)" rf"", re.DOTALL, ) _PARAM_PATTERN = re.compile( rf"<{re.escape(DSML_TOKEN)}parameter\s+name=\"([^\"]+)\"\s+string=\"(true|false)\">" rf"(.*?)" rf"", re.DOTALL, ) def parse_dsml_output(text: str) -> list[ToolCallItem] | None: """Parse DSML function_calls block from model output text. Args: text: The text containing the DSML function_calls block (including the start/end markers). Returns: List of ToolCallItem, or None if parsing fails. """ tool_calls: list[ToolCallItem] = [] for invoke_match in _INVOKE_PATTERN.finditer(text): func_name = invoke_match.group(1) invoke_body = invoke_match.group(2) args: dict[str, Any] = {} for param_match in _PARAM_PATTERN.finditer(invoke_body): param_name = param_match.group(1) is_string = param_match.group(2) == "true" param_value = param_match.group(3) if is_string: args[param_name] = param_value else: try: args[param_name] = json.loads(param_value) except (json.JSONDecodeError, ValueError): args[param_name] = param_value tool_calls.append( ToolCallItem( name=func_name, arguments=json.dumps(args), ) ) return tool_calls if tool_calls else None ================================================ FILE: src/exo/worker/engines/mlx/generator/__init__.py ================================================ ================================================ FILE: src/exo/worker/engines/mlx/generator/batch_generate.py ================================================ import time from dataclasses import dataclass, field from typing import Callable, cast import mlx.core as mx from mlx_lm.generate import ( BatchGenerator as MlxBatchGenerator, ) from mlx_lm.models.cache import RotatingKVCache from mlx_lm.sample_utils import make_logits_processors, make_sampler from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper from exo.api.types import ( CompletionTokensDetails, FinishReason, GenerationStats, PromptTokensDetails, TopLogprobItem, Usage, ) from exo.shared.types.memory import Memory from exo.shared.types.mlx import KVCacheType, Model from exo.shared.types.text_generation import TextGenerationTaskParams from exo.shared.types.worker.runner_response import GenerationResponse from exo.worker.engines.mlx.cache import ( CacheSnapshot, KVPrefixCache, encode_prompt, make_kv_cache, ) from exo.worker.engines.mlx.constants import DEFAULT_TOP_LOGPROBS, MAX_TOKENS from exo.worker.engines.mlx.generator.generate import ( ban_token_ids, eos_ids_from_tokenizer, extract_top_logprobs, prefill, ) from exo.worker.engines.mlx.utils_mlx import fix_unmatched_think_end_tokens from exo.worker.runner.bootstrap import logger _MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5 def _stop_sequences(task_params: TextGenerationTaskParams) -> list[str]: if task_params.stop is None: return [] if isinstance(task_params.stop, str): return [task_params.stop] return task_params.stop @dataclass class _EngineTask: uid: int task_params: TextGenerationTaskParams all_prompt_tokens: mx.array prefix_hit_length: int matched_index: int | None cache_snapshots: list[CacheSnapshot] | None detokenizer: StreamingDetokenizer on_generation_token: Callable[[], None] | None = None generated_text_parts: list[str] = field(default_factory=list) potential_stop_sequence_text: str = "" completion_tokens: int = 0 generation_start_time: float = 0.0 in_thinking: bool = False reasoning_tokens: int = 0 prefill_tps: float = 0.0 @dataclass(eq=False) class ExoBatchGenerator: model: Model tokenizer: TokenizerWrapper group: mx.distributed.Group | None kv_prefix_cache: KVPrefixCache | None _exo_gen: MlxBatchGenerator = field(init=False) _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False) def __post_init__(self) -> None: self._exo_gen = MlxBatchGenerator( model=self.model, stop_tokens=set(eos_ids_from_tokenizer(self.tokenizer)), prefill_step_size=4096, ) @property def has_work(self) -> bool: return ( bool(self._active_tasks) or bool(self._exo_gen.unprocessed_prompts) or self._exo_gen.active_batch is not None ) def submit( self, task_params: TextGenerationTaskParams, prompt: str, on_prefill_progress: Callable[[int, int], None] | None = None, distributed_prompt_progress_callback: Callable[[], None] | None = None, on_generation_token: Callable[[], None] | None = None, ) -> int: all_prompt_tokens = encode_prompt(self.tokenizer, prompt) all_prompt_tokens = fix_unmatched_think_end_tokens( all_prompt_tokens, self.tokenizer ) is_bench = task_params.bench prefix_hit_length = 0 matched_index: int | None = None prompt_tokens = all_prompt_tokens if self.kv_prefix_cache is not None and not is_bench: cache, remaining_tokens, matched_index = self.kv_prefix_cache.get_kv_cache( self.model, all_prompt_tokens ) prefix_hit_length = len(all_prompt_tokens) - len(remaining_tokens) if prefix_hit_length > 0: logger.info( f"KV cache hit: {prefix_hit_length}/{len(all_prompt_tokens)} tokens " f"cached ({100 * prefix_hit_length / len(all_prompt_tokens):.1f}%)" ) prompt_tokens = remaining_tokens else: cache = make_kv_cache(self.model) else: cache = make_kv_cache(self.model) seed = task_params.seed if task_params.seed is not None else 42 mx.random.seed(seed) sampler = make_sampler( temp=task_params.temperature if task_params.temperature is not None else 0.7, top_p=task_params.top_p if task_params.top_p is not None else 1.0, min_p=task_params.min_p if task_params.min_p is not None else 0.05, top_k=task_params.top_k if task_params.top_k is not None else 0, ) _prefill_tps, _prefill_tokens, cache_snapshots = prefill( self.model, self.tokenizer, sampler, prompt_tokens[:-1], cache, self.group, on_prefill_progress, distributed_prompt_progress_callback, ) # We need to clamp rotating kv caches to max size so that mlx lm's _merge_caches behaves for c in cache: if ( isinstance(c, RotatingKVCache) and c.keys is not None and c.values is not None and c.keys.shape[2] > c.max_size ): trim_size = c.keys.shape[2] - c.max_size c.keys = c._trim(trim_size, c.keys) c.values = c._trim(trim_size, c.values) c._idx = c.max_size if not is_bench: self._save_prefix_cache( all_prompt_tokens, list(cache), cache_snapshots, prefix_hit_length, matched_index, ) last_tokens = prompt_tokens[-2:] logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = ( make_logits_processors( repetition_penalty=task_params.repetition_penalty, repetition_context_size=task_params.repetition_context_size, ) ) if is_bench: # Only sample length eos tokens eos_ids = eos_ids_from_tokenizer(self.tokenizer) logits_processors = [ban_token_ids(eos_ids)] + logits_processors max_tokens = task_params.max_output_tokens or MAX_TOKENS uids = self._exo_gen.insert( prompts=[last_tokens.tolist()], max_tokens=[max_tokens], caches=[list(cache)], samplers=[sampler], logits_processors=[logits_processors], ) assert len(uids) == 1 uid = uids[0] self._active_tasks[uid] = _EngineTask( uid=uid, task_params=task_params, all_prompt_tokens=all_prompt_tokens, prefix_hit_length=prefix_hit_length, matched_index=matched_index, cache_snapshots=cache_snapshots or None, detokenizer=self.tokenizer.detokenizer, on_generation_token=on_generation_token, generation_start_time=time.perf_counter(), prefill_tps=_prefill_tps, ) return uid def step(self) -> list[tuple[int, GenerationResponse]]: if not self.has_work: return [] responses = self._exo_gen.next() results: list[tuple[int, GenerationResponse]] = [] for response in responses: if response.uid not in self._active_tasks: logger.warning( f"response uid {response.uid} was not found - should be active" ) continue state = self._active_tasks[response.uid] if state.on_generation_token is not None: state.on_generation_token() if response.finish_reason != "stop": state.detokenizer.add_token(response.token) if response.finish_reason is not None: state.detokenizer.finalize() text = state.detokenizer.last_segment state.completion_tokens += 1 state.generated_text_parts.append(text) state.potential_stop_sequence_text += text think_start = self.tokenizer.think_start think_end = self.tokenizer.think_end if think_start is not None and text == think_start: state.in_thinking = True elif think_end is not None and text == think_end: state.in_thinking = False if state.in_thinking: state.reasoning_tokens += 1 finish_reason: FinishReason | None = cast( FinishReason | None, response.finish_reason ) task_params = state.task_params stop_sequences = _stop_sequences(task_params) max_stop_len = max((len(s) for s in stop_sequences), default=0) if stop_sequences: for stop_seq in stop_sequences: if stop_seq in state.potential_stop_sequence_text: stop_index = state.potential_stop_sequence_text.find(stop_seq) text_before_stop = state.potential_stop_sequence_text[ :stop_index ] chunk_start = len(state.potential_stop_sequence_text) - len( text ) text = text_before_stop[chunk_start:] finish_reason = "stop" break is_done = finish_reason is not None logprob: float | None = None top_logprobs: list[TopLogprobItem] | None = None if task_params.logprobs: logprob, top_logprobs = extract_top_logprobs( logprobs=response.logprobs, tokenizer=self.tokenizer, top_logprobs=task_params.top_logprobs or DEFAULT_TOP_LOGPROBS, selected_token=response.token, ) stats: GenerationStats | None = None usage: Usage | None = None if is_done: try: mlx_stats = self._exo_gen.stats() generation_tps = mlx_stats.generation_tps except ZeroDivisionError: generation_elapsed = ( time.perf_counter() - state.generation_start_time ) generation_tps = ( state.completion_tokens / generation_elapsed if generation_elapsed > 0 else 0.0 ) stats = GenerationStats( prompt_tps=state.prefill_tps, generation_tps=generation_tps, prompt_tokens=len(state.all_prompt_tokens), generation_tokens=state.completion_tokens, peak_memory_usage=Memory.from_gb(mx.get_peak_memory() / 1e9), ) total_prompt_tokens = len(state.all_prompt_tokens) usage = Usage( prompt_tokens=total_prompt_tokens, completion_tokens=state.completion_tokens, total_tokens=total_prompt_tokens + state.completion_tokens, prompt_tokens_details=PromptTokensDetails( cached_tokens=state.prefix_hit_length ), completion_tokens_details=CompletionTokensDetails( reasoning_tokens=state.reasoning_tokens ), ) results.append( ( response.uid, GenerationResponse( text=text, token=response.token, logprob=logprob, top_logprobs=top_logprobs, finish_reason=finish_reason, stats=stats, usage=usage, ), ) ) if is_done: del self._active_tasks[response.uid] elif ( max_stop_len > 0 and len(state.potential_stop_sequence_text) > max_stop_len ): state.potential_stop_sequence_text = state.potential_stop_sequence_text[ -max_stop_len: ] return results def cancel(self, uids: list[int]) -> None: self._exo_gen.remove(uids) for uid in uids: self._active_tasks.pop(uid, None) def close(self) -> None: self._exo_gen.close() def _save_prefix_cache( self, all_prompt_tokens: mx.array, cache: KVCacheType, cache_snapshots: list[CacheSnapshot] | None, prefix_hit_length: int, matched_index: int | None, ) -> None: if self.kv_prefix_cache is None: return try: hit_ratio = ( prefix_hit_length / len(all_prompt_tokens) if len(all_prompt_tokens) > 0 else 0.0 ) if ( matched_index is not None and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE ): self.kv_prefix_cache.update_kv_cache( matched_index, all_prompt_tokens, cache, cache_snapshots, restore_pos=prefix_hit_length, ) else: self.kv_prefix_cache.add_kv_cache( all_prompt_tokens, cache, cache_snapshots ) except Exception: logger.warning("Failed to save prefix cache", exc_info=True) ================================================ FILE: src/exo/worker/engines/mlx/generator/generate.py ================================================ import functools import math import time from copy import deepcopy from typing import Callable, Generator, cast, get_args import mlx.core as mx from mlx_lm.generate import ( maybe_quantize_kv_cache, stream_generate, ) from mlx_lm.models.cache import ArraysCache, RotatingKVCache from mlx_lm.sample_utils import make_logits_processors, make_sampler from mlx_lm.tokenizer_utils import TokenizerWrapper from exo.api.types import ( CompletionTokensDetails, FinishReason, GenerationStats, PromptTokensDetails, TopLogprobItem, Usage, ) from exo.shared.types.common import ModelId from exo.shared.types.memory import Memory from exo.shared.types.mlx import KVCacheType, Model from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams from exo.shared.types.worker.runner_response import ( GenerationResponse, ) from exo.worker.engines.mlx.auto_parallel import ( PipelineFirstLayer, PipelineLastLayer, clear_prefill_sends, flush_prefill_sends, set_pipeline_prefill, set_pipeline_queue_sends, ) from exo.worker.engines.mlx.cache import ( CacheSnapshot, KVPrefixCache, encode_prompt, has_non_kv_caches, make_kv_cache, snapshot_ssm_states, ) from exo.worker.engines.mlx.constants import ( DEFAULT_TOP_LOGPROBS, KV_BITS, KV_GROUP_SIZE, MAX_TOKENS, ) from exo.worker.engines.mlx.utils_mlx import ( apply_chat_template, fix_unmatched_think_end_tokens, mx_barrier, ) from exo.worker.runner.bootstrap import logger generation_stream = mx.new_stream(mx.default_device()) _MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5 class PrefillCancelled(BaseException): """Raised when prefill is cancelled via the progress callback.""" def _has_pipeline_communication_layer(model: Model): for layer in model.layers: if isinstance(layer, (PipelineFirstLayer, PipelineLastLayer)): return True return False def pipeline_parallel_prefill( model: Model, prompt: mx.array, prompt_cache: KVCacheType, prefill_step_size: int, kv_group_size: int | None, kv_bits: int | None, prompt_progress_callback: Callable[[int, int], None], distributed_prompt_progress_callback: Callable[[], None] | None, group: mx.distributed.Group, ) -> None: """Prefill the KV cache for pipeline parallel with overlapping stages. Each rank processes the full prompt through its real cache, offset by leading and trailing dummy iterations. Total iterations per rank = N_real_chunks + world_size - 1: - rank r leading dummies (skip_pipeline_io, throwaway cache) - N_real_chunks real (pipeline IO active, real cache) - (world_size-1-r) trailing dummies (skip_pipeline_io, throwaway cache) e.g. Timeline (2 ranks, 3 chunks of 10240 tokens @ step=4096): iter 0: R0 real[0:4096] R1 dummy iter 1: R0 real[4096:8192] R1 real[0:4096] iter 2: R0 real[8192:10240] R1 real[4096:8192] iter 3: R0 dummy R1 real[8192:10240] This function is designed to match mlx_lm's stream_generate exactly in terms of side effects (given the same prefill step size) """ prefill_step_size = prefill_step_size // min(4, group.size()) quantize_cache_fn: Callable[..., None] = functools.partial( maybe_quantize_kv_cache, quantized_kv_start=0, kv_group_size=kv_group_size, kv_bits=kv_bits, ) _prompt_cache: KVCacheType = prompt_cache rank = group.rank() world_size = group.size() # Build list of real prompt chunk sizes total = len(prompt) real_chunk_sizes: list[int] = [] remaining = total - 1 while remaining: n = min(prefill_step_size, remaining) real_chunk_sizes.append(n) remaining -= n n_real = len(real_chunk_sizes) # Each rank does: [rank leading dummies] [N real chunks] [world_size-1-rank trailing dummies] n_leading = rank n_trailing = world_size - 1 - rank n_total = n_leading + n_real + n_trailing t_start = time.perf_counter() processed = 0 logger.info( f"[R{rank}] Pipeline prefill: {n_real} real + {n_leading} leading + {n_trailing} trailing = {n_total} iterations" ) clear_prefill_sends() # Initial callback matching generate_step prompt_progress_callback(0, total) try: with mx.stream(generation_stream): for _ in range(n_leading): if distributed_prompt_progress_callback is not None: distributed_prompt_progress_callback() for i in range(n_real): chunk_size = real_chunk_sizes[i] model( prompt[processed : processed + chunk_size][None], cache=_prompt_cache, ) quantize_cache_fn(_prompt_cache) processed += chunk_size if distributed_prompt_progress_callback is not None: distributed_prompt_progress_callback() flush_prefill_sends() prompt_progress_callback(processed, total) for _ in range(n_trailing): if distributed_prompt_progress_callback is not None: distributed_prompt_progress_callback() finally: clear_prefill_sends() # Post-loop: process remaining 1 token + add +1 entry to match stream_generate. for _ in range(2): with mx.stream(generation_stream): model(prompt[-1:][None], cache=_prompt_cache) quantize_cache_fn(_prompt_cache) flush_prefill_sends() assert _prompt_cache is not None mx.eval([c.state for c in _prompt_cache]) # type: ignore # Final callback matching generate_step prompt_progress_callback(total, total) logger.info( f"[R{rank}] Prefill: {n_real} real + {n_leading}+{n_trailing} dummy iterations, " f"Processed {processed} tokens in {(time.perf_counter() - t_start) * 1000:.1f}ms" ) def prefill( model: Model, tokenizer: TokenizerWrapper, sampler: Callable[[mx.array], mx.array], prompt_tokens: mx.array, cache: KVCacheType, group: mx.distributed.Group | None, on_prefill_progress: Callable[[int, int], None] | None, distributed_prompt_progress_callback: Callable[[], None] | None, ) -> tuple[float, int, list[CacheSnapshot]]: """Prefill the KV cache with prompt tokens. This runs the model over the prompt tokens to populate the cache, then trims off the extra generated token. Returns: (tokens_per_sec, num_tokens, snapshots) """ num_tokens = len(prompt_tokens) if num_tokens == 0: return 0.0, 0, [] logger.debug(f"Prefilling {num_tokens} tokens...") start_time = time.perf_counter() has_ssm = has_non_kv_caches(cache) snapshots: list[CacheSnapshot] = [] # TODO(evan): kill the callbacks/runner refactor def progress_callback(processed: int, total: int) -> None: elapsed = time.perf_counter() - start_time tok_per_sec = processed / elapsed if elapsed > 0 else 0 logger.debug( f"Prefill progress: {processed}/{total} tokens ({tok_per_sec:.1f} tok/s)" ) if has_ssm: snapshots.append(snapshot_ssm_states(cache)) if on_prefill_progress is not None: on_prefill_progress(processed, total) def combined_progress_callback(processed: int, total: int) -> None: if distributed_prompt_progress_callback is not None: distributed_prompt_progress_callback() progress_callback(processed, total) set_pipeline_prefill(model, is_prefill=True) mx_barrier(group) logger.info("Starting prefill") is_pipeline = _has_pipeline_communication_layer(model) prefill_step_size = 4096 try: if is_pipeline and num_tokens >= prefill_step_size: set_pipeline_queue_sends(model, queue_sends=True) assert group is not None, "Pipeline prefill requires a distributed group" pipeline_parallel_prefill( model=model, prompt=prompt_tokens, prompt_cache=cache, prefill_step_size=prefill_step_size, kv_group_size=KV_GROUP_SIZE, kv_bits=KV_BITS, prompt_progress_callback=progress_callback, distributed_prompt_progress_callback=distributed_prompt_progress_callback, group=group, ) else: # Use max_tokens=1 because max_tokens=0 does not work. # We just throw away the generated token - we only care about filling the cache for _ in stream_generate( model=model, tokenizer=tokenizer, prompt=prompt_tokens, max_tokens=1, sampler=sampler, prompt_cache=cache, prefill_step_size=prefill_step_size, kv_group_size=KV_GROUP_SIZE, kv_bits=KV_BITS, prompt_progress_callback=combined_progress_callback, ): break # Stop after first iteration - cache is now filled except PrefillCancelled: set_pipeline_queue_sends(model, queue_sends=False) set_pipeline_prefill(model, is_prefill=False) raise set_pipeline_queue_sends(model, queue_sends=False) set_pipeline_prefill(model, is_prefill=False) # stream_generate added 1 extra generated token to the cache, so we should trim it. # Because of needing to roll back arrays cache, we will generate on 2 tokens so trim 1 more. pre_gen = deepcopy(snapshots[-2]) if has_ssm else None for i, c in enumerate(cache): if has_ssm and isinstance(c, (ArraysCache, RotatingKVCache)): assert pre_gen is not None if pre_gen.states[i] is not None: cache[i] = deepcopy(pre_gen.states[i]) # type: ignore else: assert not isinstance(c, (ArraysCache, RotatingKVCache)) c.trim(2) elapsed = time.perf_counter() - start_time tokens_per_sec = num_tokens / elapsed if elapsed > 0 else 0.0 logger.debug( f"Prefill complete: {num_tokens} tokens in {elapsed:.2f}s " f"({tokens_per_sec:.1f} tok/s)" ) # Exclude the last snapshot return tokens_per_sec, num_tokens, snapshots[:-1] if snapshots else [] def warmup_inference( model: Model, tokenizer: TokenizerWrapper, group: mx.distributed.Group | None, model_id: ModelId, ) -> int: logger.info(f"warming up inference for instance: {model_id}") t = time.monotonic() content = "Prompt to warm up the inference engine. Repeat this." warmup_prompt = apply_chat_template( tokenizer=tokenizer, task_params=TextGenerationTaskParams( model=ModelId(""), input=[InputMessage(role="user", content=content)], ), ) tokens_generated = 0 cache = make_kv_cache( model=model, ) # Use a default sampler for warmup sampler = make_sampler(temp=0.0) mx_barrier(group) logger.info("Generating warmup tokens") for _r in stream_generate( model=model, tokenizer=tokenizer, prompt=warmup_prompt, max_tokens=50, sampler=sampler, prompt_cache=cache, prefill_step_size=2048, kv_group_size=KV_GROUP_SIZE, kv_bits=KV_BITS, ): logger.info("Generated warmup token: " + str(_r.text)) tokens_generated += 1 logger.info("Generated ALL warmup tokens") mx_barrier(group) logger.info(f"warmed up by generating {tokens_generated} tokens") check_for_cancel_every = min( math.ceil(tokens_generated / min(time.monotonic() - t, 0.001)), 100 ) if group is not None: check_for_cancel_every = int( mx.max( mx.distributed.all_gather( mx.array([check_for_cancel_every]), group=group, ) ).item() ) logger.info( f"runner checking for cancellation every {check_for_cancel_every} tokens" ) return check_for_cancel_every def ban_token_ids(token_ids: list[int]) -> Callable[[mx.array, mx.array], mx.array]: token_ids = [int(t) for t in token_ids] def proc(_history: mx.array, logits: mx.array) -> mx.array: for tid in token_ids: logits[..., tid] = -1e9 return logits return proc def eos_ids_from_tokenizer(tokenizer: TokenizerWrapper) -> list[int]: eos: list[int] | None = getattr(tokenizer, "eos_token_ids", None) if eos is None: return [] return eos def extract_top_logprobs( logprobs: mx.array, tokenizer: TokenizerWrapper, top_logprobs: int, selected_token: int, ) -> tuple[float, list[TopLogprobItem]]: """Extract the selected token's logprob and top alternative tokens. Args: logprobs: Full vocabulary logprobs array from MLX tokenizer: Tokenizer for decoding token IDs to strings top_logprobs: Number of top alternatives to return selected_token: The token ID that was actually sampled Returns: Tuple of (selected_token_logprob, list of TopLogprobItem for top alternatives) """ # Get the logprob of the selected token selected_logprob = float(logprobs[selected_token].item()) # Get top indices (most probable tokens) # mx.argpartition gives indices that would partition the array # We negate logprobs since argpartition finds smallest, and we want largest top_logprobs = min(top_logprobs, logprobs.shape[0]) # Don't exceed vocab size top_indices = mx.argpartition(-logprobs, top_logprobs)[:top_logprobs] # Get the actual logprob values for these indices top_values = logprobs[top_indices] # Sort by logprob (descending) for consistent ordering sort_order = mx.argsort(-top_values) top_indices = top_indices[sort_order] top_values = top_values[sort_order] # Convert to list of TopLogprobItem top_logprob_items: list[TopLogprobItem] = [] for i in range(top_logprobs): token_id = int(top_indices[i].item()) token_logprob = float(top_values[i].item()) if math.isnan(token_logprob): continue # Decode token ID to string token_str = tokenizer.decode([token_id]) # Get byte representation token_bytes = list(token_str.encode("utf-8")) top_logprob_items.append( TopLogprobItem( token=token_str, logprob=token_logprob, bytes=token_bytes, ) ) return selected_logprob, top_logprob_items def mlx_generate( model: Model, tokenizer: TokenizerWrapper, task: TextGenerationTaskParams, prompt: str, kv_prefix_cache: KVPrefixCache | None, group: mx.distributed.Group | None, on_prefill_progress: Callable[[int, int], None] | None = None, distributed_prompt_progress_callback: Callable[[], None] | None = None, on_generation_token: Callable[[], None] | None = None, ) -> Generator[GenerationResponse]: # Ensure that generation stats only contains peak memory for this generation mx.reset_peak_memory() # TODO: Randomise task seed and set in taskparams, instead of hard coding as 42. seed = task.seed or 42 mx.random.seed(seed) # Encode prompt once at the top and fix unmatched think tags all_prompt_tokens = encode_prompt(tokenizer, prompt) all_prompt_tokens = fix_unmatched_think_end_tokens(all_prompt_tokens, tokenizer) # Do not use the prefix cache if we are trying to do benchmarks. is_bench = task.bench if is_bench: kv_prefix_cache = None # Use prefix cache if available, otherwise create fresh cache prefix_hit_length = 0 matched_index: int | None = None if kv_prefix_cache is None: caches = make_kv_cache(model=model) prompt_tokens = all_prompt_tokens else: caches, prompt_tokens, matched_index = kv_prefix_cache.get_kv_cache( model, all_prompt_tokens ) prefix_hit_length = len(all_prompt_tokens) - len(prompt_tokens) if prefix_hit_length > 0: logger.info( f"KV cache hit: {prefix_hit_length}/{len(all_prompt_tokens)} tokens cached ({100 * prefix_hit_length / len(all_prompt_tokens):.1f}%)" ) logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = ( make_logits_processors( repetition_penalty=task.repetition_penalty, repetition_context_size=task.repetition_context_size, ) ) if is_bench: # Only sample length eos tokens eos_ids = eos_ids_from_tokenizer(tokenizer) logits_processors = [ban_token_ids(eos_ids)] + logits_processors sampler = make_sampler( temp=task.temperature if task.temperature is not None else 0.7, top_p=task.top_p if task.top_p is not None else 1.0, min_p=task.min_p if task.min_p is not None else 0.05, top_k=task.top_k if task.top_k is not None else 0, ) # Normalize stop sequences to a list stop_sequences: list[str] = ( ([task.stop] if isinstance(task.stop, str) else task.stop) if task.stop is not None else [] ) max_stop_len = max((len(s) for s in stop_sequences), default=0) # Prefill cache with all tokens except the last one prefill_tps, prefill_tokens, ssm_snapshots_list = prefill( model, tokenizer, sampler, prompt_tokens[:-1], caches, group, on_prefill_progress, distributed_prompt_progress_callback, ) cache_snapshots: list[CacheSnapshot] | None = ssm_snapshots_list or None # stream_generate starts from the last token last_token = prompt_tokens[-2:] max_tokens = task.max_output_tokens or MAX_TOKENS accumulated_text = "" generated_text_parts: list[str] = [] generation_start_time = time.perf_counter() usage: Usage | None = None in_thinking = False reasoning_tokens = 0 think_start = tokenizer.think_start think_end = tokenizer.think_end logger.info("Starting decode") mx_barrier(group) for completion_tokens, out in enumerate( stream_generate( model=model, tokenizer=tokenizer, prompt=last_token, max_tokens=max_tokens, sampler=sampler, logits_processors=logits_processors, prompt_cache=caches, prefill_step_size=1, kv_group_size=KV_GROUP_SIZE, kv_bits=KV_BITS, ), start=1, ): generated_text_parts.append(out.text) accumulated_text += out.text if think_start is not None and out.text == think_start: in_thinking = True elif think_end is not None and out.text == think_end: in_thinking = False if in_thinking: reasoning_tokens += 1 # Check for stop sequences text = out.text finish_reason: FinishReason | None = cast( FinishReason | None, out.finish_reason ) stop_matched = False if stop_sequences: for stop_seq in stop_sequences: if stop_seq in accumulated_text: # Trim text to just before the stop sequence stop_index = accumulated_text.find(stop_seq) text_before_stop = accumulated_text[:stop_index] chunk_start = len(accumulated_text) - len(out.text) text = text_before_stop[chunk_start:] finish_reason = "stop" stop_matched = True break is_done = finish_reason is not None stats: GenerationStats | None = None if is_done: stats = GenerationStats( prompt_tps=float(prefill_tps or out.prompt_tps), generation_tps=float(out.generation_tps), prompt_tokens=int(prefill_tokens + out.prompt_tokens), generation_tokens=int(out.generation_tokens), peak_memory_usage=Memory.from_gb(out.peak_memory), ) if not stop_matched and out.finish_reason not in get_args(FinishReason): logger.warning( f"Model generated unexpected finish_reason: {out.finish_reason}" ) total_prompt_tokens = len(all_prompt_tokens) usage = Usage( prompt_tokens=total_prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_prompt_tokens + completion_tokens, prompt_tokens_details=PromptTokensDetails( cached_tokens=prefix_hit_length ), completion_tokens_details=CompletionTokensDetails( reasoning_tokens=reasoning_tokens ), ) # Extract logprobs from the full vocabulary logprobs array logprob: float | None = None top_logprobs: list[TopLogprobItem] | None = None if task.logprobs: logprob, top_logprobs = extract_top_logprobs( logprobs=out.logprobs, tokenizer=tokenizer, top_logprobs=task.top_logprobs or DEFAULT_TOP_LOGPROBS, selected_token=out.token, ) if is_done: # Log generation stats generation_elapsed = time.perf_counter() - generation_start_time generated_tokens = len(generated_text_parts) generation_tps = ( generated_tokens / generation_elapsed if generation_elapsed > 0 else 0.0 ) logger.debug( f"Generation complete: prefill {prompt_tokens} tokens @ " f"{prefill_tps:.1f} tok/s, generated {generated_tokens} tokens @ " f"{generation_tps:.1f} tok/s" ) if kv_prefix_cache is not None: generated_tokens_array = mx.array( tokenizer.encode( "".join(generated_text_parts), add_special_tokens=False ) ) full_prompt_tokens = mx.concatenate( [all_prompt_tokens, generated_tokens_array] ) hit_ratio = ( prefix_hit_length / len(all_prompt_tokens) if len(all_prompt_tokens) > 0 else 0.0 ) if ( matched_index is not None and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE ): kv_prefix_cache.update_kv_cache( matched_index, full_prompt_tokens, caches, cache_snapshots, restore_pos=prefix_hit_length, ) else: kv_prefix_cache.add_kv_cache( full_prompt_tokens, caches, cache_snapshots ) if on_generation_token is not None: on_generation_token() yield GenerationResponse( text=text, token=out.token, logprob=logprob, top_logprobs=top_logprobs, finish_reason=finish_reason, stats=stats, usage=usage, ) if is_done: mx_barrier(group) break # Limit accumulated_text to what's needed for stop sequence detection if max_stop_len > 0 and len(accumulated_text) > max_stop_len: accumulated_text = accumulated_text[-max_stop_len:] ================================================ FILE: src/exo/worker/engines/mlx/tests/test_batch_generate.py ================================================ # pyright: reportAny=false, reportUnknownVariableType=false # pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false # pyright: reportUnknownLambdaType=false, reportPrivateUsage=false # pyright: reportInvalidCast=false, reportArgumentType=false # pyright: reportUnusedImport=false """Test B=1 vs B=2 equivalence for batch generation. Verifies that running two requests concurrently in a batch (B=2) produces identical token selections to running them sequentially (B=1). Uses random weights — no model download required. """ from pathlib import Path from typing import cast import mlx.core as mx import mlx.nn as nn import mlx.utils import pytest from mlx_lm.generate import _merge_caches from mlx_lm.sample_utils import make_sampler from mlx_lm.tokenizer_utils import TokenizerWrapper from transformers import AutoTokenizer # Import batch_generate to activate the right-padding BatchKVCache patch import exo.worker.engines.mlx.generator.batch_generate # noqa: F401 from exo.shared.types.mlx import Model from exo.worker.engines.mlx.cache import encode_prompt, make_kv_cache from exo.worker.engines.mlx.generator.generate import prefill NUM_STEPS = 20 def _init_random(model: nn.Module) -> None: """Initialize all model parameters with random values.""" params = model.parameters() new_params = mlx.utils.tree_map( lambda p: mx.random.normal(shape=p.shape, dtype=p.dtype) if isinstance(p, mx.array) else p, params, ) model.update(new_params) mx.eval(model.parameters()) def _run_b1_vs_b2( model: Model, tokenizer: TokenizerWrapper, tokens_a: mx.array, tokens_b: mx.array, ) -> tuple[float, int]: """Run B=1 sequential and B=2 batched, return (max_diff, mismatches).""" sampler = make_sampler(temp=0.0) # B=1 sequential cache_a1 = make_kv_cache(model) prefill(model, tokenizer, sampler, tokens_a[:-1], cache_a1, None, None, None) merged_a1 = _merge_caches([[c for c in cache_a1]]) for c in merged_a1: c.prepare(lengths=[1], right_padding=[0]) model(mx.array([[tokens_a[-2].item()]]), cache=merged_a1) mx.eval([c.state for c in merged_a1]) for c in merged_a1: c.finalize() cache_b1 = make_kv_cache(model) prefill(model, tokenizer, sampler, tokens_b[:-1], cache_b1, None, None, None) merged_b1 = _merge_caches([[c for c in cache_b1]]) for c in merged_b1: c.prepare(lengths=[1], right_padding=[0]) model(mx.array([[tokens_b[-2].item()]]), cache=merged_b1) mx.eval([c.state for c in merged_b1]) for c in merged_b1: c.finalize() b1_logits_a: list[mx.array] = [] b1_logits_b: list[mx.array] = [] next_a, next_b = tokens_a[-1].item(), tokens_b[-1].item() for _ in range(NUM_STEPS): la = model(mx.array([[next_a]]), cache=merged_a1) mx.eval(la) b1_logits_a.append(la[0, -1]) next_a = int(mx.argmax(la[0, -1]).item()) lb = model(mx.array([[next_b]]), cache=merged_b1) mx.eval(lb) b1_logits_b.append(lb[0, -1]) next_b = int(mx.argmax(lb[0, -1]).item()) # B=2 batched cache_a2 = make_kv_cache(model) cache_b2 = make_kv_cache(model) prefill(model, tokenizer, sampler, tokens_a[:-1], cache_a2, None, None, None) prefill(model, tokenizer, sampler, tokens_b[:-1], cache_b2, None, None, None) merged_b2 = _merge_caches([list(cache_a2), list(cache_b2)]) for c in merged_b2: c.prepare(lengths=[1, 1], right_padding=[0, 0]) model( mx.array([[tokens_a[-2].item()], [tokens_b[-2].item()]]), cache=merged_b2, ) mx.eval([c.state for c in merged_b2]) for c in merged_b2: c.finalize() b2_logits_a: list[mx.array] = [] b2_logits_b: list[mx.array] = [] next_a2, next_b2 = tokens_a[-1].item(), tokens_b[-1].item() for _ in range(NUM_STEPS): l2 = model(mx.array([[next_a2], [next_b2]]), cache=merged_b2) mx.eval(l2) b2_logits_a.append(l2[0, -1]) b2_logits_b.append(l2[1, -1]) next_a2 = int(mx.argmax(l2[0, -1]).item()) next_b2 = int(mx.argmax(l2[1, -1]).item()) # Compare max_diff = 0.0 mismatches = 0 for step in range(NUM_STEPS): diff_a = float( mx.max( mx.abs( b1_logits_a[step].astype(mx.float32) - b2_logits_a[step].astype(mx.float32) ) ).item() ) diff_b = float( mx.max( mx.abs( b1_logits_b[step].astype(mx.float32) - b2_logits_b[step].astype(mx.float32) ) ).item() ) max_diff = max(max_diff, diff_a, diff_b) if int(mx.argmax(b1_logits_a[step]).item()) != int( mx.argmax(b2_logits_a[step]).item() ): mismatches += 1 if int(mx.argmax(b1_logits_b[step]).item()) != int( mx.argmax(b2_logits_b[step]).item() ): mismatches += 1 return max_diff, mismatches def _make_tokenizer() -> TokenizerWrapper: """Load the Qwen tokenizer (tiny download, shared across Qwen models).""" from huggingface_hub import snapshot_download model_path = Path( snapshot_download( "mlx-community/Qwen3.5-35B-A3B-4bit", allow_patterns=["tokenizer*", "*.jinja"], ) ) hf_tokenizer = AutoTokenizer.from_pretrained(model_path) return TokenizerWrapper(hf_tokenizer) @pytest.mark.slow def test_batch_b2_llama() -> None: """Llama-style model (KVCache only) must produce bit-exact logits in B=2. Right-padded BatchKVCache keeps data at position 0 for all sequences, so flash attention sees identical data layout as B=1 → bit-exact output. """ from mlx_lm.models.llama import Model as LlamaModel from mlx_lm.models.llama import ModelArgs mx.random.seed(42) args = ModelArgs( model_type="llama", hidden_size=256, num_hidden_layers=4, intermediate_size=512, num_attention_heads=4, num_key_value_heads=2, rms_norm_eps=1e-6, vocab_size=248320, rope_theta=10000.0, tie_word_embeddings=True, ) model = LlamaModel(args) _init_random(model) tokenizer = _make_tokenizer() tokens_a = encode_prompt(tokenizer, "Write a short essay about AI.") tokens_b = encode_prompt(tokenizer, "Explain evolution briefly.") max_diff, mismatches = _run_b1_vs_b2( cast(Model, model), tokenizer, tokens_a, tokens_b ) assert mismatches == 0, f"Llama B=2 token mismatches: {mismatches}/{NUM_STEPS * 2}" assert max_diff < 0.002, f"Llama B=2 max logit diff: {max_diff}" @pytest.mark.slow def test_batch_b2_qwen35_moe() -> None: """Qwen3.5 MoE model (hybrid SSM+attention+MoE) must produce bit-exact logits in B=2. Right-padded BatchKVCache keeps data at position 0 for all sequences, so flash attention sees identical data layout as B=1 → bit-exact output. """ from mlx_lm.models.qwen3_5_moe import Model as Qwen35MoeModel from mlx_lm.models.qwen3_5_moe import ModelArgs mx.random.seed(42) config = { "model_type": "qwen3_5_moe", "text_config": { "model_type": "qwen3_5_moe_text", "hidden_size": 256, "num_hidden_layers": 8, "intermediate_size": 512, "num_attention_heads": 4, "num_key_value_heads": 2, "rms_norm_eps": 1e-6, "vocab_size": 248320, "head_dim": 64, "max_position_embeddings": 4096, "full_attention_interval": 4, "layer_types": [ "linear_attention", "linear_attention", "linear_attention", "full_attention", "linear_attention", "linear_attention", "linear_attention", "full_attention", ], "linear_conv_kernel_dim": 4, "linear_key_head_dim": 64, "linear_num_key_heads": 4, "linear_num_value_heads": 4, "linear_value_head_dim": 64, "mamba_ssm_dtype": "float32", "num_experts": 8, "num_experts_per_tok": 2, "moe_intermediate_size": 256, "shared_expert_intermediate_size": 256, "rope_parameters": { "rope_type": "default", "rope_theta": 10000000, }, "attention_bias": False, "attn_output_gate": True, }, } args = ModelArgs.from_dict(config) model = Qwen35MoeModel(args) _init_random(model) tokenizer = _make_tokenizer() tokens_a = encode_prompt(tokenizer, "Write a short essay about AI.") tokens_b = encode_prompt(tokenizer, "Explain evolution briefly.") max_diff, mismatches = _run_b1_vs_b2( cast(Model, model), tokenizer, tokens_a, tokens_b ) assert mismatches == 0, ( f"Qwen3.5 MoE B=2 token mismatches: {mismatches}/{NUM_STEPS * 2}" ) assert max_diff < 0.002, f"Qwen3.5 MoE B=2 max logit diff: {max_diff}" ================================================ FILE: src/exo/worker/engines/mlx/utils_mlx.py ================================================ import json import os import re import sys import tempfile import time from pathlib import Path from typing import Any, cast # Monkey-patch for transformers 5.x compatibility # Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location # which was moved in transformers 5.0.0rc2 try: import transformers.models.gpt2.tokenization_gpt2 as gpt2_tokenization from transformers.convert_slow_tokenizer import bytes_to_unicode if not hasattr(gpt2_tokenization, "bytes_to_unicode"): gpt2_tokenization.bytes_to_unicode = bytes_to_unicode # type: ignore[attr-defined] except ImportError: pass # transformers < 5.0 or bytes_to_unicode not available from mlx_lm.models.cache import KVCache from mlx_lm.models.deepseek_v3 import DeepseekV3Model from mlx_lm.tokenizer_utils import TokenizerWrapper from exo.shared.models.model_cards import ModelId from exo.worker.engines.mlx.constants import TRUST_REMOTE_CODE try: from mlx_lm.tokenizer_utils import load_tokenizer except ImportError: from mlx_lm.tokenizer_utils import load as load_tokenizer import contextlib import mlx.core as mx import mlx.nn as nn from mlx_lm.utils import load_model from pydantic import RootModel from exo.download.download_utils import build_model_path from exo.shared.types.common import Host from exo.shared.types.memory import Memory from exo.shared.types.mlx import Model from exo.shared.types.tasks import TaskId, TextGeneration from exo.shared.types.text_generation import TextGenerationTaskParams from exo.shared.types.worker.instances import ( BoundInstance, MlxJacclInstance, MlxRingInstance, ) from exo.shared.types.worker.shards import ( CfgShardMetadata, PipelineShardMetadata, ShardMetadata, TensorShardMetadata, ) from exo.worker.engines.mlx.auto_parallel import ( LayerLoadedCallback, TimeoutCallback, eval_with_timeout, get_inner_model, get_layers, pipeline_auto_parallel, tensor_auto_parallel, ) from exo.worker.runner.bootstrap import logger Group = mx.distributed.Group def get_weights_size(model_shard_meta: ShardMetadata) -> Memory: return Memory.from_float_kb( (model_shard_meta.end_layer - model_shard_meta.start_layer) / model_shard_meta.n_layers * model_shard_meta.model_card.storage_size.in_kb / ( 1 if isinstance(model_shard_meta, PipelineShardMetadata) else model_shard_meta.world_size ) ) class ModelLoadingTimeoutError(Exception): pass class HostList(RootModel[list[str]]): @classmethod def from_hosts(cls, hosts: list[Host]) -> "HostList": return cls(root=[str(host) for host in hosts]) def mlx_distributed_init( bound_instance: BoundInstance, ) -> Group: """ Initialize MLX distributed. """ rank = bound_instance.bound_shard.device_rank logger.info(f"Starting initialization for rank {rank}") with tempfile.TemporaryDirectory() as tmpdir: coordination_file = str( Path(tmpdir) / f"hosts_{bound_instance.instance.instance_id}_{rank}.json" ) # TODO: singleton instances match bound_instance.instance: case MlxRingInstance(hosts_by_node=hosts_by_node, ephemeral_port=_): hosts_for_node = hosts_by_node[bound_instance.bound_node_id] hosts_json = HostList.from_hosts(hosts_for_node).model_dump_json() with open(coordination_file, "w") as f: _ = f.write(hosts_json) logger.info( f"rank {rank} hostfile: {coordination_file} hosts: {hosts_json}" ) os.environ["MLX_HOSTFILE"] = coordination_file os.environ["MLX_RANK"] = str(rank) os.environ["MLX_RING_VERBOSE"] = "1" group = mx.distributed.init(backend="ring", strict=True) case MlxJacclInstance( jaccl_devices=jaccl_devices, jaccl_coordinators=jaccl_coordinators ): assert all( jaccl_devices[i][i] is None for i in range(len(jaccl_devices)) ) # Use RDMA connectivity matrix jaccl_devices_json = json.dumps(jaccl_devices) with open(coordination_file, "w") as f: _ = f.write(jaccl_devices_json) jaccl_coordinator = jaccl_coordinators[bound_instance.bound_node_id] logger.info( f"rank {rank} MLX_IBV_DEVICES: {coordination_file} with devices: {jaccl_devices_json}" ) logger.info(f"rank {rank} MLX_JACCL_COORDINATOR: {jaccl_coordinator}") os.environ["MLX_IBV_DEVICES"] = coordination_file os.environ["MLX_RANK"] = str(rank) os.environ["MLX_JACCL_COORDINATOR"] = jaccl_coordinator group = mx.distributed.init(backend="jaccl", strict=True) logger.info(f"Rank {rank} mlx distributed initialization complete") return group def initialize_mlx( bound_instance: BoundInstance, ) -> Group: # should we unseed it? # TODO: pass in seed from params mx.random.seed(42) assert len(bound_instance.instance.shard_assignments.node_to_runner) > 1, ( "Tried to initialize mlx for a single node instance" ) return mlx_distributed_init(bound_instance) def load_mlx_items( bound_instance: BoundInstance, group: Group | None, on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, ) -> tuple[Model, TokenizerWrapper]: if group is None: logger.info(f"Single device used for {bound_instance.instance}") model_path = build_model_path(bound_instance.bound_shard.model_card.model_id) start_time = time.perf_counter() model, _ = load_model(model_path, lazy=True, strict=False) # Eval layers one by one for progress reporting try: inner = get_inner_model(model) layers = get_layers(inner) total = len(layers) for i, layer in enumerate(layers): mx.eval(layer) # type: ignore if on_layer_loaded is not None: on_layer_loaded(i, total) except ValueError as e: logger.opt(exception=e).debug( "Model architecture doesn't support layer-by-layer progress tracking", ) mx.eval(model) end_time = time.perf_counter() logger.info(f"Time taken to load model: {(end_time - start_time):.2f}s") tokenizer = get_tokenizer(model_path, bound_instance.bound_shard) else: logger.info("Starting distributed init") start_time = time.perf_counter() model, tokenizer = shard_and_load( bound_instance.bound_shard, group=group, on_timeout=on_timeout, on_layer_loaded=on_layer_loaded, ) end_time = time.perf_counter() logger.info( f"Time taken to shard and load model: {(end_time - start_time):.2f}s" ) set_wired_limit_for_model(get_weights_size(bound_instance.bound_shard)) mx.clear_cache() return cast(Model, model), tokenizer def shard_and_load( shard_metadata: ShardMetadata, group: Group, on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, ) -> tuple[nn.Module, TokenizerWrapper]: model_path = build_model_path(shard_metadata.model_card.model_id) model, _ = load_model(model_path, lazy=True, strict=False) logger.debug(model) if hasattr(model, "model") and isinstance(model.model, DeepseekV3Model): # type: ignore pass # TODO: See if we should quantize the model. # def is_attention_layer(path: str) -> bool: # path = path.lower() # return "self_attn" in path and "layernorm" not in path # def quant_predicate(path: str, module: nn.Module): # if not isinstance(module, nn.Linear): # return False # return is_attention_layer(path) # model, config = quantize_model( # model, config, group_size=KV_GROUP_SIZE, bits=ATTENTION_KV_BITS, quant_predicate=quant_predicate, mode=QUANTIZE_MODEL_MODE # ) assert isinstance(model, nn.Module) tokenizer = get_tokenizer(model_path, shard_metadata) logger.info(f"Group size: {group.size()}, group rank: {group.rank()}") # Estimate timeout based on model size (5x default for large queued workloads) base_timeout = float(os.environ.get("EXO_MODEL_LOAD_TIMEOUT", "300")) model_size = get_weights_size(shard_metadata) timeout_seconds = base_timeout + model_size.in_gb logger.info( f"Evaluating model parameters with timeout of {timeout_seconds:.0f}s " f"(model size: {model_size.in_gb:.1f}GB)" ) match shard_metadata: case TensorShardMetadata(): logger.info(f"loading model from {model_path} with tensor parallelism") model = tensor_auto_parallel( model, group, timeout_seconds, on_timeout, on_layer_loaded ) case PipelineShardMetadata(): logger.info(f"loading model from {model_path} with pipeline parallelism") model = pipeline_auto_parallel( model, group, shard_metadata, on_layer_loaded=on_layer_loaded ) eval_with_timeout(model.parameters(), timeout_seconds, on_timeout) case CfgShardMetadata(): raise ValueError( "CfgShardMetadata is not supported for text model loading - " "this metadata type is only for image generation models" ) # TODO: Do we need this? mx.eval(model) logger.debug("SHARDED") logger.debug(model) # Synchronize processes before generation to avoid timeout mx_barrier(group) return model, tokenizer def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata) -> TokenizerWrapper: """Load tokenizer for a model shard. Delegates to load_tokenizer_for_model_id.""" return load_tokenizer_for_model_id( shard_metadata.model_card.model_id, model_path, trust_remote_code=shard_metadata.model_card.trust_remote_code, ) def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None: """ Get the EOS token IDs for a model based on its ID. Some models require explicit EOS token configuration that isn't in their tokenizer config. This function returns the known EOS token IDs for such models. Args: model_id: The HuggingFace model ID Returns: List of EOS token IDs, or None if the model uses standard tokenizer config """ model_id_lower = model_id.lower() if "kimi-k2" in model_id_lower: return [163586] elif "glm-5" in model_id_lower or "glm-4.7" in model_id_lower: # For GLM-5 and GLM-4.7 # 154820: <|endoftext|>, 154827: <|user|>, 154829: <|observation|> return [154820, 154827, 154829] elif "glm" in model_id_lower: # For GLM-4.5 and older return [151336, 151329, 151338] elif "gpt-oss" in model_id_lower: return [200002, 200012] elif "qwen3.5" in model_id_lower or "qwen-3.5" in model_id_lower: # For Qwen3.5: 248046 (<|im_end|>), 248044 (<|endoftext|>) return [248046, 248044] return None def load_tokenizer_for_model_id( model_id: ModelId, model_path: Path, *, trust_remote_code: bool = TRUST_REMOTE_CODE ) -> TokenizerWrapper: """ Load tokenizer for a model given its ID and local path. This is the core tokenizer loading logic, handling special cases for different model families (Kimi, GLM, etc.) and transformers 5.x compatibility. Args: model_id: The HuggingFace model ID (e.g., "moonshotai/Kimi-K2-Instruct") model_path: Local path where the model/tokenizer files are stored Returns: TokenizerWrapper instance configured for the model """ model_id_lower = model_id.lower() eos_token_ids = get_eos_token_ids_for_model(model_id) # Kimi uses a custom TikTokenTokenizer that transformers 5.x can't load via AutoTokenizer if "kimi-k2" in model_id_lower: import importlib.util import types sys.path.insert(0, str(model_path)) # Load tool_declaration_ts first (tokenization_kimi imports it with relative import) tool_decl_path = model_path / "tool_declaration_ts.py" if tool_decl_path.exists(): spec = importlib.util.spec_from_file_location( "tool_declaration_ts", tool_decl_path ) if spec and spec.loader: tool_decl_module = importlib.util.module_from_spec(spec) sys.modules["tool_declaration_ts"] = tool_decl_module spec.loader.exec_module(tool_decl_module) # Load tokenization_kimi with patched source (convert relative to absolute import) tok_path = model_path / "tokenization_kimi.py" source = tok_path.read_text() source = source.replace("from .tool_declaration_ts", "from tool_declaration_ts") spec = importlib.util.spec_from_file_location("tokenization_kimi", tok_path) if spec: tok_module = types.ModuleType("tokenization_kimi") tok_module.__file__ = str(tok_path) sys.modules["tokenization_kimi"] = tok_module exec(compile(source, tok_path, "exec"), tok_module.__dict__) # noqa: S102 TikTokenTokenizer = tok_module.TikTokenTokenizer # type: ignore[attr-defined] # noqa: N806 else: from tokenization_kimi import TikTokenTokenizer # type: ignore[import-not-found] # noqa: I001 hf_tokenizer: Any = TikTokenTokenizer.from_pretrained(model_path) # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType] # Patch encode to use internal tiktoken model directly # transformers 5.x has a bug in the encode->pad path for slow tokenizers def _patched_encode(text: str, **_kwargs: object) -> list[int]: # Pass allowed_special="all" to handle special tokens like <|im_user|> return list(hf_tokenizer.model.encode(text, allowed_special="all")) # pyright: ignore[reportUnknownMemberType,reportUnknownArgumentType] hf_tokenizer.encode = _patched_encode return TokenizerWrapper( hf_tokenizer, eos_token_ids=eos_token_ids, tool_call_start="<|tool_calls_section_begin|>", tool_call_end="<|tool_calls_section_end|>", tool_parser=_parse_kimi_tool_calls, ) tokenizer = load_tokenizer( model_path, tokenizer_config_extra={"trust_remote_code": trust_remote_code}, eos_token_ids=eos_token_ids, ) if "gemma-3" in model_id_lower: gemma_3_eos_id = 1 gemma_3_end_of_turn_id = 106 if tokenizer.eos_token_ids is not None: if gemma_3_end_of_turn_id not in tokenizer.eos_token_ids: tokenizer.eos_token_ids = list(tokenizer.eos_token_ids) + [ gemma_3_end_of_turn_id ] else: tokenizer.eos_token_ids = [gemma_3_eos_id, gemma_3_end_of_turn_id] return tokenizer def _normalize_tool_calls(msg_dict: dict[str, Any]) -> None: """Normalize tool_calls in a message dict. OpenAI format has tool_calls[].function.arguments as a JSON string, but some chat templates (e.g., GLM) expect it as a dict. """ tool_calls = msg_dict.get("tool_calls") if not tool_calls or not isinstance(tool_calls, list): return for tc in tool_calls: # pyright: ignore[reportUnknownVariableType] if not isinstance(tc, dict): continue func = tc.get("function") # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] if not isinstance(func, dict): continue args = func.get("arguments") # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] if isinstance(args, str): with contextlib.suppress(json.JSONDecodeError): func["arguments"] = json.loads(args) def _collect_nested_property_names(schema: dict[str, Any]) -> set[str]: names: set[str] = set() properties: dict[str, Any] = schema.get("properties", {}) # type: ignore[reportAny] for prop_spec in properties.values(): # pyright: ignore[reportAny] if not isinstance(prop_spec, dict): continue if prop_spec.get("type") == "array": # type: ignore[reportAny] items: dict[str, Any] | None = prop_spec.get("items") # type: ignore[reportAny] if isinstance(items, dict) and items.get("type") == "object": # type: ignore[reportAny] inner_props: dict[str, Any] = items.get("properties", {}) # type: ignore[reportAny] for k in inner_props: # pyright: ignore[reportUnknownVariableType] names.add(str(k)) # pyright: ignore[reportUnknownArgumentType] names.update(_collect_nested_property_names(items)) # pyright: ignore[reportUnknownArgumentType] return names def _schemas_lost_in_prompt(prompt: str, tools: list[dict[str, Any]]) -> bool: """Return True if nested property names from any tool schema are absent.""" for tool in tools: fn: dict[str, Any] = tool.get("function", {}) # type: ignore params: dict[str, Any] = fn.get("parameters", {}) # type: ignore nested = _collect_nested_property_names(params) if nested and not all(name in prompt for name in nested): return True return False _LOSSY_TEMPLATE_PATTERN = re.compile( r"""inner_type\s*==\s*["']object \| object["']\s*or\s*inner_type\|length\s*>\s*\d+""", ) def _patch_lossy_chat_template(template: str) -> str | None: """Patch chat templates that collapse nested object schemas to ``any[]``. Some templates (e.g., GPT-OSS) have a guard like:: inner_type == "object | object" or inner_type|length > 50 The length check silently drops complex array-of-object schemas. We remove the length guard, keeping only the object-union check. Returns the patched template, or *None* if no patch was needed. """ patched, n = _LOSSY_TEMPLATE_PATTERN.subn( lambda m: m.group(0).split(" or ")[0], # keep only the object-union check template, ) return patched if n > 0 else None def _needs_dsml_encoding(task_params: TextGenerationTaskParams) -> bool: if "deepseek-v3.2" not in task_params.model.lower(): return False # Use DSML encoding when tools are provided or tool results are in the conversation if task_params.tools: return True if task_params.chat_template_messages: return any( msg.get("role") == "tool" for msg in task_params.chat_template_messages ) return False def apply_chat_template( tokenizer: TokenizerWrapper, task_params: TextGenerationTaskParams, ) -> str: """Convert TextGenerationTaskParams to a chat template prompt. Converts the internal format (input + instructions) to a messages list that can be processed by the tokenizer's chat template. When chat_template_messages is available (from Chat Completions API), uses those directly to preserve tool_calls, thinking, and other fields. """ formatted_messages: list[dict[str, Any]] = [] if task_params.chat_template_messages is not None: # Use pre-formatted messages that preserve tool_calls, thinking, etc. formatted_messages = list(task_params.chat_template_messages) for msg in formatted_messages: _normalize_tool_calls(msg) else: # Add system message (instructions) if present if task_params.instructions: formatted_messages.append( {"role": "system", "content": task_params.instructions} ) # Convert input to messages for msg in task_params.input: if not msg.content: logger.warning("Received message with empty content, skipping") continue formatted_messages.append({"role": msg.role, "content": msg.content}) # For assistant prefilling, append content after templating to avoid a closing turn token. partial_assistant_content: str | None = None if formatted_messages and formatted_messages[-1].get("role") == "assistant": partial_assistant_content = cast(str, formatted_messages[-1].get("content", "")) formatted_messages = formatted_messages[:-1] if _needs_dsml_encoding(task_params): from exo.worker.engines.mlx.dsml_encoding import encode_messages prompt = encode_messages( messages=formatted_messages, thinking_mode="thinking" if task_params.enable_thinking else "chat", tools=task_params.tools, ) if partial_assistant_content: prompt += partial_assistant_content logger.info(prompt) return prompt extra_kwargs: dict[str, Any] = {} if task_params.enable_thinking is not None: # Qwen3 and GLM use "enable_thinking"; DeepSeek uses "thinking". # Jinja ignores unknown variables, so passing both is safe. extra_kwargs["enable_thinking"] = task_params.enable_thinking extra_kwargs["thinking"] = task_params.enable_thinking if task_params.reasoning_effort is not None: extra_kwargs["reasoning_effort"] = task_params.reasoning_effort patched_template: str | None = None if task_params.tools: original_template: str | None = getattr(tokenizer, "chat_template", None) if isinstance(original_template, str): patched_template = _patch_lossy_chat_template(original_template) if patched_template is not None: logger.info( "Patched lossy chat template (removed inner_type length guard)" ) prompt: str = tokenizer.apply_chat_template( formatted_messages, tokenize=False, add_generation_prompt=True, tools=task_params.tools, **({"chat_template": patched_template} if patched_template is not None else {}), **extra_kwargs, ) if task_params.tools and _schemas_lost_in_prompt(prompt, task_params.tools): logger.warning("Chat template lost nested tool schemas even after patching") if partial_assistant_content: prompt += partial_assistant_content logger.info(prompt) return prompt def detect_thinking_prompt_suffix(prompt: str, tokenizer: TokenizerWrapper) -> bool: """ Detect if prompt ends with a thinking opening tag that should be prepended to the output stream. """ think_token = tokenizer.think_start return think_token is not None and prompt.rstrip().endswith(think_token) def fix_unmatched_think_end_tokens( tokens: mx.array, tokenizer: TokenizerWrapper ) -> mx.array: if not tokenizer.has_thinking: return tokens assert tokenizer.think_start_id assert tokenizer.think_end_id think_start_id: int = tokenizer.think_start_id think_end_id: int = tokenizer.think_end_id token_list: list[int] = cast(list[int], tokens.tolist()) result: list[int] = [] depth = 0 for token in token_list: if token == think_start_id: depth += 1 elif token == think_end_id: if depth == 0: result.append(think_start_id) else: depth -= 1 result.append(token) return mx.array(result) class NullKVCache(KVCache): """ A KVCache that pretends to exist but holds zero tokens. It satisfies .state/.meta_state and never allocates real keys/values. """ def __init__(self, dtype: mx.Dtype = mx.float16): super().__init__() # zero-length K/V so shapes/dtypes are defined but empty self.keys = mx.zeros((1, 1, 0, 1), dtype=dtype) self.values = mx.zeros((1, 1, 0, 1), dtype=dtype) self.offset = 0 @property def state(self) -> tuple[mx.array, mx.array]: # matches what mx.save_safetensors / mx.eval expect return self.keys, self.values @state.setter def state(self, v: tuple[mx.array, mx.array]) -> None: raise NotImplementedError("We should not be setting a NullKVCache.") def mlx_force_oom(size: int = 200000) -> None: """ Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations. """ mx.set_default_device(mx.gpu) a = mx.random.uniform(shape=(size, size), dtype=mx.float32) b = mx.random.uniform(shape=(size, size), dtype=mx.float32) mx.eval(a, b) c = mx.matmul(a, b) d = mx.matmul(a, c) e = mx.matmul(b, c) f = mx.sigmoid(d + e) mx.eval(f) def set_wired_limit_for_model(model_size: Memory): """ A context manager to temporarily change the wired limit. Note, the wired limit should not be changed during an async eval. If an async eval could be running pass in the streams to synchronize with prior to exiting the context manager. """ if not mx.metal.is_available(): return max_rec_size = Memory.from_bytes( int(mx.device_info()["max_recommended_working_set_size"]) ) if model_size > 0.9 * max_rec_size: logger.warning( f"Generating with a model that requires {model_size.in_float_mb:.1f} MB " f"which is close to the maximum recommended size of {max_rec_size.in_float_mb:.1f} " "MB. This can be slow. See the documentation for possible work-arounds: " "https://github.com/ml-explore/mlx-lm/tree/main#large-models" ) mx.set_wired_limit(max_rec_size.in_bytes) logger.info(f"Wired limit set to {max_rec_size}.") def mlx_cleanup( model: Model | None, tokenizer: TokenizerWrapper | None, group: Group | None ) -> None: del model, tokenizer, group mx.clear_cache() import gc gc.collect() def mx_any(bool_: bool, group: Group | None) -> bool: if group is None: return bool_ num_true = mx.distributed.all_sum( mx.array(bool_), group=group, stream=mx.default_stream(mx.Device(mx.cpu)) ) mx.eval(num_true) return num_true.item() > 0 def mx_barrier(group: Group | None): if group is None: return mx.eval( mx.distributed.all_sum( mx.array(1.0), group=group, stream=mx.default_stream(mx.Device(mx.cpu)) ) ) def _parse_kimi_tool_calls(text: str): import regex as re # kimi has a fixed function naming scheme, with a json formatted arg # functions.multiply:0<|tool_call_argument_begin|>{"a": 2, "b": 3} _func_name_regex = re.compile( r"^\s*((?:functions\.)?(.+?):\d+)\s*<\|tool_call_argument_begin\|>", re.DOTALL ) _func_arg_regex = re.compile(r"<\|tool_call_argument_begin\|>\s*(.*)\s*", re.DOTALL) _tool_call_split_regex = re.compile( r"<\|tool_call_begin\|>(.*?)<\|tool_call_end\|>", re.DOTALL ) def _parse_single_tool(text: str) -> dict[str, Any]: func_name_match = _func_name_regex.search(text) if func_name_match is None: raise ValueError("No tool call found.") tool_call_id = func_name_match.group(1) # e.g. "functions.get_weather:0" func_name = func_name_match.group(2) # e.g. "get_weather" func_args_match = _func_arg_regex.search(text) if func_args_match is None: raise ValueError("No tool call arguments found.") func_args = func_args_match.group(1) try: arg_dct = json.loads(func_args) # pyright: ignore[reportAny] except Exception: arg_dct = None return dict(id=tool_call_id, name=func_name, arguments=arg_dct) tool_matches = _tool_call_split_regex.findall(text) if tool_matches: return [_parse_single_tool(match) for match in tool_matches] # pyright: ignore[reportAny] else: return [_parse_single_tool(text)] def mx_all_gather_tasks( tasks: list[TextGeneration], group: mx.distributed.Group | None, ) -> tuple[list[TextGeneration], list[TextGeneration]]: def encode_task_id(task_id: TaskId) -> list[int]: utf8_task_id = task_id.encode() return [ int.from_bytes(utf8_task_id[i : i + 1]) for i in range(len(utf8_task_id)) ] def decode_task_id(encoded_task_id: list[int]) -> TaskId: return TaskId( bytes.decode(b"".join((x).to_bytes(length=1) for x in encoded_task_id)) ) uuid_byte_length = 36 n_tasks = len(tasks) all_counts = cast( list[int], mx.distributed.all_gather(mx.array([n_tasks]), group=group).tolist(), ) max_tasks = max(all_counts) world_size: int = 1 if group is None else group.size() if max_tasks == 0: return [], [] padded = [encode_task_id(task.task_id) for task in tasks] + [ [0] * uuid_byte_length ] * (max_tasks - n_tasks) assert all(len(encoded_task_id) == uuid_byte_length for encoded_task_id in padded) gathered = cast( list[list[list[int]]], mx.distributed.all_gather(mx.array(padded), group=group) .reshape(world_size, max_tasks, -1) .tolist(), ) all_task_ids: list[list[TaskId]] = [ [decode_task_id(encoded_task_id) for encoded_task_id in rank_tasks[:count]] for rank_tasks, count in zip(gathered, all_counts, strict=True) ] agreed_ids = set[TaskId].intersection(*(set(tids) for tids in all_task_ids)) local_tasks = {task.task_id: task for task in tasks} agreed = [local_tasks[tid] for tid in sorted(agreed_ids)] different = [task for task in tasks if task.task_id not in agreed_ids] return agreed, different ================================================ FILE: src/exo/worker/main.py ================================================ from collections import defaultdict from datetime import datetime, timezone import anyio from anyio import fail_after from loguru import logger from exo.api.types import ImageEditsTaskParams from exo.download.download_utils import resolve_model_in_path from exo.shared.apply import apply from exo.shared.models.model_cards import ModelId from exo.shared.types.commands import ( ForwarderCommand, ForwarderDownloadCommand, StartDownload, ) from exo.shared.types.common import CommandId, NodeId, SystemId from exo.shared.types.events import ( Event, IndexedEvent, InputChunkReceived, NodeDownloadProgress, NodeGatheredInfo, TaskCreated, TaskStatusUpdated, TopologyEdgeCreated, TopologyEdgeDeleted, ) from exo.shared.types.multiaddr import Multiaddr from exo.shared.types.state import State from exo.shared.types.tasks import ( CancelTask, CreateRunner, DownloadModel, ImageEdits, Shutdown, Task, TaskStatus, ) from exo.shared.types.topology import Connection, SocketConnection from exo.shared.types.worker.downloads import DownloadCompleted from exo.shared.types.worker.runners import RunnerId from exo.utils.channels import Receiver, Sender, channel from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer from exo.utils.info_gatherer.net_profile import check_reachable from exo.utils.keyed_backoff import KeyedBackoff from exo.utils.task_group import TaskGroup from exo.worker.plan import plan from exo.worker.runner.runner_supervisor import RunnerSupervisor class Worker: def __init__( self, node_id: NodeId, *, event_receiver: Receiver[IndexedEvent], event_sender: Sender[Event], # This is for requesting updates. It doesn't need to be a general command sender right now, # but I think it's the correct way to be thinking about commands command_sender: Sender[ForwarderCommand], download_command_sender: Sender[ForwarderDownloadCommand], ): self.node_id: NodeId = node_id self.event_receiver = event_receiver self.event_sender = event_sender self.command_sender = command_sender self.download_command_sender = download_command_sender self.state: State = State() self.runners: dict[RunnerId, RunnerSupervisor] = {} self._tg: TaskGroup = TaskGroup() self._system_id = SystemId() # Buffer for input image chunks (for image editing) self.input_chunk_buffer: dict[CommandId, dict[int, str]] = {} self.input_chunk_counts: dict[CommandId, int] = {} self._download_backoff: KeyedBackoff[ModelId] = KeyedBackoff(base=0.5, cap=10.0) async def run(self): logger.info("Starting Worker") info_send, info_recv = channel[GatheredInfo]() info_gatherer: InfoGatherer = InfoGatherer(info_send) try: async with self._tg as tg: tg.start_soon(info_gatherer.run) tg.start_soon(self._forward_info, info_recv) tg.start_soon(self.plan_step) tg.start_soon(self._event_applier) tg.start_soon(self._poll_connection_updates) finally: # Actual shutdown code - waits for all tasks to complete before executing. logger.info("Stopping Worker") self.event_sender.close() self.command_sender.close() self.download_command_sender.close() for runner in self.runners.values(): runner.shutdown() async def _forward_info(self, recv: Receiver[GatheredInfo]): with recv as info_stream: async for info in info_stream: await self.event_sender.send( NodeGatheredInfo( node_id=self.node_id, when=str(datetime.now(tz=timezone.utc)), info=info, ) ) async def _event_applier(self): with self.event_receiver as events: async for event in events: # 2. for each event, apply it to the state self.state = apply(self.state, event=event) event = event.event # Buffer input image chunks for image editing if isinstance(event, InputChunkReceived): cmd_id = event.command_id if cmd_id not in self.input_chunk_buffer: self.input_chunk_buffer[cmd_id] = {} self.input_chunk_counts[cmd_id] = event.chunk.total_chunks self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = ( event.chunk.data ) async def plan_step(self): while True: await anyio.sleep(0.1) task: Task | None = plan( self.node_id, self.runners, self.state.downloads, self.state.instances, self.state.runners, self.state.tasks, self.input_chunk_buffer, self.input_chunk_counts, ) if task is None: continue # Gate DownloadModel on backoff BEFORE emitting TaskCreated # to prevent flooding the event log with useless events if isinstance(task, DownloadModel): model_id = task.shard_metadata.model_card.model_id if not self._download_backoff.should_proceed(model_id): continue logger.info(f"Worker plan: {task.__class__.__name__}") assert task.task_status await self.event_sender.send(TaskCreated(task_id=task.task_id, task=task)) # lets not kill the worker if a runner is unresponsive match task: case CreateRunner(): self._create_supervisor(task) await self.event_sender.send( TaskStatusUpdated( task_id=task.task_id, task_status=TaskStatus.Complete ) ) case DownloadModel(shard_metadata=shard): model_id = shard.model_card.model_id self._download_backoff.record_attempt(model_id) found_path = resolve_model_in_path(model_id) if found_path is not None: logger.info( f"Model {model_id} found in EXO_MODELS_PATH at {found_path}" ) await self.event_sender.send( NodeDownloadProgress( download_progress=DownloadCompleted( node_id=self.node_id, shard_metadata=shard, model_directory=str(found_path), total=shard.model_card.storage_size, read_only=True, ) ) ) await self.event_sender.send( TaskStatusUpdated( task_id=task.task_id, task_status=TaskStatus.Complete, ) ) else: await self.download_command_sender.send( ForwarderDownloadCommand( origin=self._system_id, command=StartDownload( target_node_id=self.node_id, shard_metadata=shard, ), ) ) await self.event_sender.send( TaskStatusUpdated( task_id=task.task_id, task_status=TaskStatus.Running, ) ) case Shutdown(runner_id=runner_id): runner = self.runners.pop(runner_id) try: with fail_after(3): await runner.start_task(task) except TimeoutError: await self.event_sender.send( TaskStatusUpdated( task_id=task.task_id, task_status=TaskStatus.TimedOut ) ) finally: runner.shutdown() case CancelTask( cancelled_task_id=cancelled_task_id, runner_id=runner_id ): await self.runners[runner_id].cancel_task(cancelled_task_id) await self.event_sender.send( TaskStatusUpdated( task_id=task.task_id, task_status=TaskStatus.Complete ) ) case ImageEdits() if task.task_params.total_input_chunks > 0: # Assemble image from chunks and inject into task cmd_id = task.command_id chunks = self.input_chunk_buffer.get(cmd_id, {}) assembled = "".join(chunks[i] for i in range(len(chunks))) logger.info( f"Assembled input image from {len(chunks)} chunks, " f"total size: {len(assembled)} bytes" ) # Create modified task with assembled image data modified_task = ImageEdits( task_id=task.task_id, command_id=task.command_id, instance_id=task.instance_id, task_status=task.task_status, task_params=ImageEditsTaskParams( image_data=assembled, total_input_chunks=task.task_params.total_input_chunks, prompt=task.task_params.prompt, model=task.task_params.model, n=task.task_params.n, quality=task.task_params.quality, output_format=task.task_params.output_format, response_format=task.task_params.response_format, size=task.task_params.size, image_strength=task.task_params.image_strength, bench=task.task_params.bench, stream=task.task_params.stream, partial_images=task.task_params.partial_images, advanced_params=task.task_params.advanced_params, ), ) # Cleanup buffers if cmd_id in self.input_chunk_buffer: del self.input_chunk_buffer[cmd_id] if cmd_id in self.input_chunk_counts: del self.input_chunk_counts[cmd_id] await self._start_runner_task(modified_task) case task: await self._start_runner_task(task) def shutdown(self): self._tg.cancel_tasks() async def _start_runner_task(self, task: Task): if (instance := self.state.instances.get(task.instance_id)) is not None: await self.runners[ instance.shard_assignments.node_to_runner[self.node_id] ].start_task(task) def _create_supervisor(self, task: CreateRunner) -> RunnerSupervisor: """Creates and stores a new AssignedRunner with initial downloading status.""" runner = RunnerSupervisor.create( bound_instance=task.bound_instance, event_sender=self.event_sender.clone(), ) self.runners[task.bound_instance.bound_runner_id] = runner self._tg.start_soon(runner.run) return runner async def _poll_connection_updates(self): while True: edges = set( conn.edge for conn in self.state.topology.out_edges(self.node_id) ) conns: defaultdict[NodeId, set[str]] = defaultdict(set) async for ip, nid in check_reachable( self.state.topology, self.node_id, self.state.node_network, ): if ip in conns[nid]: continue conns[nid].add(ip) edge = SocketConnection( # nonsense multiaddr sink_multiaddr=Multiaddr(address=f"/ip4/{ip}/tcp/52415") if "." in ip # nonsense multiaddr else Multiaddr(address=f"/ip6/{ip}/tcp/52415"), ) if edge not in edges: logger.debug(f"ping discovered {edge=}") await self.event_sender.send( TopologyEdgeCreated( conn=Connection(source=self.node_id, sink=nid, edge=edge) ) ) for conn in self.state.topology.out_edges(self.node_id): if not isinstance(conn.edge, SocketConnection): continue # ignore mDNS discovered connections if conn.edge.sink_multiaddr.port != 52415: continue if ( conn.sink not in conns or conn.edge.sink_multiaddr.ip_address not in conns[conn.sink] ): logger.debug(f"ping failed to discover {conn=}") await self.event_sender.send(TopologyEdgeDeleted(conn=conn)) await anyio.sleep(10) ================================================ FILE: src/exo/worker/plan.py ================================================ # pyright: reportUnusedImport = false from collections.abc import Mapping, Sequence from exo.shared.types.common import CommandId, NodeId from exo.shared.types.tasks import ( CancelTask, ConnectToGroup, CreateRunner, DownloadModel, ImageEdits, ImageGeneration, LoadModel, Shutdown, StartWarmup, Task, TaskId, TaskStatus, TextGeneration, ) from exo.shared.types.worker.downloads import ( DownloadCompleted, DownloadFailed, DownloadOngoing, DownloadProgress, ) from exo.shared.types.worker.instances import BoundInstance, Instance, InstanceId from exo.shared.types.worker.runners import ( RunnerConnected, RunnerConnecting, RunnerFailed, RunnerId, RunnerIdle, RunnerLoaded, RunnerLoading, RunnerReady, RunnerRunning, RunnerStatus, RunnerWarmingUp, ) from exo.worker.runner.runner_supervisor import RunnerSupervisor def plan( node_id: NodeId, # Runners is expected to be FRESH and so should not come from state runners: Mapping[RunnerId, RunnerSupervisor], global_download_status: Mapping[NodeId, Sequence[DownloadProgress]], instances: Mapping[InstanceId, Instance], all_runners: Mapping[RunnerId, RunnerStatus], # all global tasks: Mapping[TaskId, Task], input_chunk_buffer: Mapping[CommandId, dict[int, str]] | None = None, input_chunk_counts: Mapping[CommandId, int] | None = None, ) -> Task | None: # Python short circuiting OR logic should evaluate these sequentially. return ( _cancel_tasks(runners, tasks) or _kill_runner(runners, all_runners, instances) or _create_runner(node_id, runners, instances) or _model_needs_download(node_id, runners, global_download_status) or _init_distributed_backend(runners, all_runners) or _load_model(runners, all_runners, global_download_status) or _ready_to_warmup(runners, all_runners) or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer or {}) ) def _kill_runner( runners: Mapping[RunnerId, RunnerSupervisor], all_runners: Mapping[RunnerId, RunnerStatus], instances: Mapping[InstanceId, Instance], ) -> Shutdown | None: for runner in runners.values(): runner_id = runner.bound_instance.bound_runner_id if (instance_id := runner.bound_instance.instance.instance_id) not in instances: return Shutdown(instance_id=instance_id, runner_id=runner_id) for ( global_runner_id ) in runner.bound_instance.instance.shard_assignments.node_to_runner.values(): if runner_id == global_runner_id: continue if isinstance(all_runners.get(global_runner_id, None), RunnerFailed): return Shutdown( instance_id=instance_id, runner_id=runner_id, ) def _create_runner( node_id: NodeId, runners: Mapping[RunnerId, RunnerSupervisor], instances: Mapping[InstanceId, Instance], ) -> CreateRunner | None: for instance in instances.values(): runner_id = instance.shard_assignments.node_to_runner.get(node_id, None) if runner_id is None: continue if runner_id in runners: continue shard = instance.shard(runner_id) assert shard is not None return CreateRunner( instance_id=instance.instance_id, bound_instance=BoundInstance( instance=instance, bound_runner_id=runner_id, bound_node_id=node_id ), ) def _model_needs_download( node_id: NodeId, runners: Mapping[RunnerId, RunnerSupervisor], global_download_status: Mapping[NodeId, Sequence[DownloadProgress]], ) -> DownloadModel | None: local_downloads = global_download_status.get(node_id, []) download_status = { dp.shard_metadata.model_card.model_id: dp for dp in local_downloads } for runner in runners.values(): model_id = runner.bound_instance.bound_shard.model_card.model_id if isinstance(runner.status, RunnerIdle) and ( model_id not in download_status or not isinstance( download_status[model_id], (DownloadOngoing, DownloadCompleted, DownloadFailed), ) ): # We don't invalidate download_status randomly in case a file gets deleted on disk return DownloadModel( instance_id=runner.bound_instance.instance.instance_id, shard_metadata=runner.bound_instance.bound_shard, ) def _init_distributed_backend( runners: Mapping[RunnerId, RunnerSupervisor], all_runners: Mapping[RunnerId, RunnerStatus], ): for runner in runners.values(): instance = runner.bound_instance.instance shard_assignments = instance.shard_assignments is_single_node_instance = len(shard_assignments.runner_to_shard) == 1 if is_single_node_instance: continue runner_is_idle = isinstance(runner.status, RunnerIdle) all_runners_connecting = all( isinstance( all_runners.get(global_runner_id), (RunnerConnecting, RunnerIdle), ) for global_runner_id in shard_assignments.runner_to_shard ) if not (runner_is_idle and all_runners_connecting): continue runner_id = runner.bound_instance.bound_runner_id shard = runner.bound_instance.bound_shard device_rank = shard.device_rank world_size = shard.world_size assert device_rank < world_size assert device_rank >= 0 accepting_ranks = device_rank < world_size - 1 # Rank = n-1 connecting_rank_ready = device_rank == world_size - 1 and all( isinstance(all_runners.get(global_runner_id, None), RunnerConnecting) for global_runner_id in shard_assignments.runner_to_shard if global_runner_id != runner_id ) if not (accepting_ranks or connecting_rank_ready): continue return ConnectToGroup(instance_id=instance.instance_id) return None def _load_model( runners: Mapping[RunnerId, RunnerSupervisor], all_runners: Mapping[RunnerId, RunnerStatus], global_download_status: Mapping[NodeId, Sequence[DownloadProgress]], ) -> LoadModel | None: for runner in runners.values(): instance = runner.bound_instance.instance shard_assignments = instance.shard_assignments all_local_downloads_complete = all( nid in global_download_status and any( isinstance(dp, DownloadCompleted) and dp.shard_metadata.model_card.model_id == shard_assignments.model_id for dp in global_download_status[nid] ) for nid in shard_assignments.node_to_runner ) if not all_local_downloads_complete: continue is_single_node_instance = len(instance.shard_assignments.runner_to_shard) == 1 if is_single_node_instance and isinstance(runner.status, RunnerIdle): return LoadModel(instance_id=instance.instance_id) is_runner_waiting = isinstance(runner.status, RunnerConnected) all_ready_for_model = all( isinstance( all_runners.get(global_runner_id, None), (RunnerConnected, RunnerLoading, RunnerLoaded), ) for global_runner_id in shard_assignments.runner_to_shard ) if is_runner_waiting and all_ready_for_model: return LoadModel(instance_id=instance.instance_id) return None def _ready_to_warmup( runners: Mapping[RunnerId, RunnerSupervisor], all_runners: Mapping[RunnerId, RunnerStatus], ) -> StartWarmup | None: for runner in runners.values(): instance = runner.bound_instance.instance shard_assignments = instance.shard_assignments shard = runner.bound_instance.bound_shard device_rank = shard.device_rank runner_id = runner.bound_instance.bound_runner_id world_size = shard.world_size is_runner_loaded = isinstance(runner.status, RunnerLoaded) assert device_rank < world_size assert device_rank >= 0 # Rank != 0 accepting_ranks_ready = device_rank > 0 and all( isinstance( all_runners.get(global_runner_id, None), (RunnerLoaded, RunnerWarmingUp), ) for global_runner_id in shard_assignments.runner_to_shard ) # Rank = 0 connecting_rank_ready = device_rank == 0 and all( isinstance(all_runners.get(global_runner_id, None), RunnerWarmingUp) for global_runner_id in shard_assignments.runner_to_shard if global_runner_id != runner_id ) if is_runner_loaded and (accepting_ranks_ready or connecting_rank_ready): return StartWarmup(instance_id=instance.instance_id) return None def _pending_tasks( runners: Mapping[RunnerId, RunnerSupervisor], tasks: Mapping[TaskId, Task], all_runners: Mapping[RunnerId, RunnerStatus], input_chunk_buffer: Mapping[CommandId, dict[int, str]], ) -> Task | None: for task in tasks.values(): # for now, just forward chat completions # TODO(ciaran): do this better! if not isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)): continue if task.task_status not in (TaskStatus.Pending, TaskStatus.Running): continue # For ImageEdits tasks, verify all input chunks have been received if isinstance(task, ImageEdits) and task.task_params.total_input_chunks > 0: cmd_id = task.command_id expected = task.task_params.total_input_chunks received = len(input_chunk_buffer.get(cmd_id, {})) if received < expected: continue # Wait for all chunks to arrive for runner in runners.values(): if task.instance_id != runner.bound_instance.instance.instance_id: continue # the task status _should_ be set to completed by the LAST runner # it is currently set by the first # this is definitely a hack if task.task_id in runner.completed or task.task_id in runner.in_progress: continue if isinstance(runner.status, (RunnerReady, RunnerRunning)) and all( isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning)) for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard ): return task def _cancel_tasks( runners: Mapping[RunnerId, RunnerSupervisor], tasks: Mapping[TaskId, Task], ) -> Task | None: for task in tasks.values(): if task.task_status != TaskStatus.Cancelled: continue for runner_id, runner in runners.items(): if task.instance_id != runner.bound_instance.instance.instance_id: continue if task.task_id in runner.cancelled: continue return CancelTask( instance_id=task.instance_id, cancelled_task_id=task.task_id, runner_id=runner_id, ) ================================================ FILE: src/exo/worker/runner/__init__.py ================================================ ================================================ FILE: src/exo/worker/runner/bootstrap.py ================================================ import os import resource import loguru from exo.shared.types.events import Event, RunnerStatusUpdated from exo.shared.types.tasks import Task, TaskId from exo.shared.types.worker.instances import BoundInstance from exo.shared.types.worker.runners import RunnerFailed from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender logger: "loguru.Logger" = loguru.logger def entrypoint( bound_instance: BoundInstance, event_sender: MpSender[Event], task_receiver: MpReceiver[Task], cancel_receiver: MpReceiver[TaskId], _logger: "loguru.Logger", ) -> None: global logger logger = _logger soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) resource.setrlimit(resource.RLIMIT_NOFILE, (min(max(soft, 2048), hard), hard)) fast_synch_override = os.environ.get("EXO_FAST_SYNCH") if fast_synch_override != "off": os.environ["MLX_METAL_FAST_SYNCH"] = "1" else: os.environ["MLX_METAL_FAST_SYNCH"] = "0" logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}") # Import main after setting global logger - this lets us just import logger from this module try: if bound_instance.is_image_model: from exo.worker.runner.image_models.runner import Runner as ImageRunner runner = ImageRunner( bound_instance, event_sender, task_receiver, cancel_receiver ) runner.main() else: from exo.worker.runner.llm_inference.runner import Runner runner = Runner( bound_instance, event_sender, task_receiver, cancel_receiver ) runner.main() except ClosedResourceError: logger.warning("Runner communication closed unexpectedly") except Exception as e: logger.opt(exception=e).warning( f"Runner {bound_instance.bound_runner_id} crashed with critical exception {e}" ) event_sender.send( RunnerStatusUpdated( runner_id=bound_instance.bound_runner_id, runner_status=RunnerFailed(error_message=str(e)), ) ) finally: try: event_sender.close() task_receiver.close() finally: event_sender.join() task_receiver.join() logger.info("bye from the runner") ================================================ FILE: src/exo/worker/runner/image_models/__init__.py ================================================ ================================================ FILE: src/exo/worker/runner/image_models/runner.py ================================================ import base64 import time from typing import TYPE_CHECKING, Literal import mlx.core as mx from exo.api.types import ImageGenerationStats from exo.shared.constants import EXO_MAX_CHUNK_SIZE, EXO_TRACING_ENABLED from exo.shared.models.model_cards import ModelTask from exo.shared.tracing import clear_trace_buffer, get_trace_buffer from exo.shared.types.chunks import ErrorChunk, ImageChunk from exo.shared.types.common import CommandId, ModelId from exo.shared.types.events import ( ChunkGenerated, Event, RunnerStatusUpdated, TaskAcknowledged, TaskStatusUpdated, TraceEventData, TracesCollected, ) from exo.shared.types.tasks import ( CANCEL_ALL_TASKS, ConnectToGroup, ImageEdits, ImageGeneration, LoadModel, Shutdown, StartWarmup, Task, TaskId, TaskStatus, ) from exo.shared.types.worker.instances import BoundInstance from exo.shared.types.worker.runner_response import ( ImageGenerationResponse, PartialImageResponse, ) from exo.shared.types.worker.runners import ( RunnerConnected, RunnerConnecting, RunnerFailed, RunnerIdle, RunnerLoaded, RunnerLoading, RunnerReady, RunnerRunning, RunnerShutdown, RunnerShuttingDown, RunnerStatus, RunnerWarmingUp, ) from exo.shared.types.worker.shards import ( CfgShardMetadata, PipelineShardMetadata, ShardMetadata, ) from exo.utils.channels import MpReceiver, MpSender from exo.worker.engines.image import ( DistributedImageModel, generate_image, initialize_image_model, warmup_image_generator, ) from exo.worker.engines.mlx.utils_mlx import ( initialize_mlx, ) from exo.worker.runner.bootstrap import logger def _is_primary_output_node(shard_metadata: ShardMetadata) -> bool: """Check if this node is the primary output node for image generation. For CFG models: the last pipeline stage in CFG group 0 (positive prompt). For non-CFG models: the last pipeline stage. """ if isinstance(shard_metadata, CfgShardMetadata): is_pipeline_last = ( shard_metadata.pipeline_rank == shard_metadata.pipeline_world_size - 1 ) return is_pipeline_last and shard_metadata.cfg_rank == 0 elif isinstance(shard_metadata, PipelineShardMetadata): return shard_metadata.device_rank == shard_metadata.world_size - 1 return False def _process_image_response( response: ImageGenerationResponse | PartialImageResponse, command_id: CommandId, shard_metadata: ShardMetadata, event_sender: MpSender[Event], image_index: int, ) -> None: """Process a single image response and send chunks.""" encoded_data = base64.b64encode(response.image_data).decode("utf-8") is_partial = isinstance(response, PartialImageResponse) # Extract stats from final ImageGenerationResponse if available stats = response.stats if isinstance(response, ImageGenerationResponse) else None _send_image_chunk( encoded_data=encoded_data, command_id=command_id, model_id=shard_metadata.model_card.model_id, event_sender=event_sender, image_index=response.image_index, is_partial=is_partial, partial_index=response.partial_index if is_partial else None, total_partials=response.total_partials if is_partial else None, stats=stats, image_format=response.format, ) def _send_traces_if_enabled( event_sender: MpSender[Event], task_id: TaskId, rank: int, ) -> None: if not EXO_TRACING_ENABLED: return traces = get_trace_buffer() if traces: trace_data = [ TraceEventData( name=t.name, start_us=t.start_us, duration_us=t.duration_us, rank=t.rank, category=t.category, ) for t in traces ] event_sender.send( TracesCollected( task_id=task_id, rank=rank, traces=trace_data, ) ) clear_trace_buffer() def _send_image_chunk( encoded_data: str, command_id: CommandId, model_id: ModelId, event_sender: MpSender[Event], image_index: int, is_partial: bool, partial_index: int | None = None, total_partials: int | None = None, stats: ImageGenerationStats | None = None, image_format: Literal["png", "jpeg", "webp"] | None = None, ) -> None: """Send base64-encoded image data as chunks via events.""" data_chunks = [ encoded_data[i : i + EXO_MAX_CHUNK_SIZE] for i in range(0, len(encoded_data), EXO_MAX_CHUNK_SIZE) ] total_chunks = len(data_chunks) for chunk_index, chunk_data in enumerate(data_chunks): # Only include stats on the last chunk of the final image chunk_stats = ( stats if chunk_index == total_chunks - 1 and not is_partial else None ) event_sender.send( ChunkGenerated( command_id=command_id, chunk=ImageChunk( model=model_id, data=chunk_data, chunk_index=chunk_index, total_chunks=total_chunks, image_index=image_index, is_partial=is_partial, partial_index=partial_index, total_partials=total_partials, stats=chunk_stats, format=image_format, ), ) ) class Runner: def __init__( self, bound_instance: BoundInstance, event_sender: MpSender[Event], task_receiver: MpReceiver[Task], cancel_receiver: MpReceiver[TaskId], ): self.event_sender = event_sender self.task_receiver = task_receiver self.cancel_receiver = cancel_receiver self.bound_instance = bound_instance self.instance, self.runner_id, self.shard_metadata = ( bound_instance.instance, bound_instance.bound_runner_id, bound_instance.bound_shard, ) self.device_rank = self.shard_metadata.device_rank logger.info("hello from the runner") if getattr(self.shard_metadata, "immediate_exception", False): raise Exception("Fake exception - runner failed to spin up.") if timeout := getattr(self.shard_metadata, "should_timeout", 0): time.sleep(timeout) self.setup_start_time = time.time() self.cancelled_tasks = set[TaskId]() self.image_model: DistributedImageModel | None = None self.group = None self.current_status: RunnerStatus = RunnerIdle() logger.info("runner created") self.update_status(RunnerIdle()) self.seen = set[TaskId]() def update_status(self, status: RunnerStatus): self.current_status = status self.event_sender.send( RunnerStatusUpdated( runner_id=self.runner_id, runner_status=self.current_status ) ) def send_task_status(self, task: Task, status: TaskStatus): self.event_sender.send( TaskStatusUpdated(task_id=task.task_id, task_status=status) ) def acknowledge_task(self, task: Task): self.event_sender.send(TaskAcknowledged(task_id=task.task_id)) def main(self): with self.task_receiver as tasks: for task in tasks: if task.task_id in self.seen: logger.warning("repeat task - potential error") self.seen.add(task.task_id) self.cancelled_tasks.discard(CANCEL_ALL_TASKS) self.send_task_status(task, TaskStatus.Running) self.handle_task(task) was_cancelled = (task.task_id in self.cancelled_tasks) or ( CANCEL_ALL_TASKS in self.cancelled_tasks ) if not was_cancelled: self.send_task_status(task, TaskStatus.Complete) self.update_status(self.current_status) if isinstance(self.current_status, RunnerShutdown): break def handle_task(self, task: Task): match task: case ConnectToGroup() if isinstance( self.current_status, (RunnerIdle, RunnerFailed) ): logger.info("runner connecting") self.update_status(RunnerConnecting()) self.acknowledge_task(task) self.group = initialize_mlx(self.bound_instance) logger.info("runner connected") self.current_status = RunnerConnected() # we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to case LoadModel() if ( isinstance(self.current_status, RunnerConnected) and self.group is not None ) or (isinstance(self.current_status, RunnerIdle) and self.group is None): logger.info("runner loading") self.update_status(RunnerLoading()) self.acknowledge_task(task) assert ( ModelTask.TextToImage in self.shard_metadata.model_card.tasks or ModelTask.ImageToImage in self.shard_metadata.model_card.tasks ), f"Incorrect model task(s): {self.shard_metadata.model_card.tasks}" self.image_model = initialize_image_model(self.bound_instance) self.current_status = RunnerLoaded() logger.info("runner loaded") case StartWarmup() if isinstance(self.current_status, RunnerLoaded): logger.info("runner warming up") self.update_status(RunnerWarmingUp()) self.acknowledge_task(task) logger.info(f"warming up inference for instance: {self.instance}") assert self.image_model image = warmup_image_generator(model=self.image_model) if image is not None: logger.info(f"warmed up by generating {image.size} image") else: logger.info("warmup completed (non-primary node)") logger.info( f"runner initialized in {time.time() - self.setup_start_time} seconds" ) self.current_status = RunnerReady() logger.info("runner ready") case ImageGeneration(task_params=task_params, command_id=command_id) if ( isinstance(self.current_status, RunnerReady) ): assert self.image_model logger.info(f"received image generation request: {str(task)[:500]}") logger.info("runner running") self.update_status(RunnerRunning()) self.acknowledge_task(task) try: image_index = 0 for response in generate_image( model=self.image_model, task=task_params ): is_primary_output = _is_primary_output_node(self.shard_metadata) if is_primary_output: match response: case PartialImageResponse(): logger.info( f"sending partial ImageChunk {response.partial_index}/{response.total_partials}" ) _process_image_response( response, command_id, self.shard_metadata, self.event_sender, image_index, ) case ImageGenerationResponse(): logger.info("sending final ImageChunk") _process_image_response( response, command_id, self.shard_metadata, self.event_sender, image_index, ) image_index += 1 # can we make this more explicit? except Exception as e: if _is_primary_output_node(self.shard_metadata): self.event_sender.send( ChunkGenerated( command_id=command_id, chunk=ErrorChunk( model=self.shard_metadata.model_card.model_id, finish_reason="error", error_message=str(e), ), ) ) raise finally: _send_traces_if_enabled( self.event_sender, task.task_id, self.device_rank ) self.current_status = RunnerReady() logger.info("runner ready") case ImageEdits(task_params=task_params, command_id=command_id) if ( isinstance(self.current_status, RunnerReady) ): assert self.image_model logger.info(f"received image edits request: {str(task)[:500]}") logger.info("runner running") self.update_status(RunnerRunning()) self.acknowledge_task(task) try: image_index = 0 for response in generate_image( model=self.image_model, task=task_params ): if _is_primary_output_node(self.shard_metadata): match response: case PartialImageResponse(): logger.info( f"sending partial ImageChunk {response.partial_index}/{response.total_partials}" ) _process_image_response( response, command_id, self.shard_metadata, self.event_sender, image_index, ) case ImageGenerationResponse(): logger.info("sending final ImageChunk") _process_image_response( response, command_id, self.shard_metadata, self.event_sender, image_index, ) image_index += 1 except Exception as e: if _is_primary_output_node(self.shard_metadata): self.event_sender.send( ChunkGenerated( command_id=command_id, chunk=ErrorChunk( model=self.shard_metadata.model_card.model_id, finish_reason="error", error_message=str(e), ), ) ) raise finally: _send_traces_if_enabled( self.event_sender, task.task_id, self.device_rank ) self.current_status = RunnerReady() logger.info("runner ready") case Shutdown(): logger.info("runner shutting down") if not TYPE_CHECKING: del self.image_model, self.group mx.clear_cache() import gc gc.collect() self.update_status(RunnerShuttingDown()) self.acknowledge_task(task) self.current_status = RunnerShutdown() case _: raise ValueError( f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}" ) ================================================ FILE: src/exo/worker/runner/llm_inference/__init__.py ================================================ ================================================ FILE: src/exo/worker/runner/llm_inference/batch_generator.py ================================================ import itertools import time from abc import ABC, abstractmethod from collections import deque from collections.abc import Generator, Iterable from dataclasses import dataclass, field import mlx.core as mx from mlx_lm.tokenizer_utils import TokenizerWrapper from exo.shared.constants import EXO_MAX_CONCURRENT_REQUESTS from exo.shared.types.chunks import ErrorChunk, PrefillProgressChunk from exo.shared.types.common import ModelId from exo.shared.types.events import ChunkGenerated, Event from exo.shared.types.mlx import Model from exo.shared.types.tasks import CANCEL_ALL_TASKS, TaskId, TextGeneration from exo.shared.types.text_generation import TextGenerationTaskParams from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse from exo.utils.channels import MpReceiver, MpSender from exo.worker.engines.mlx.cache import KVPrefixCache from exo.worker.engines.mlx.generator.batch_generate import ExoBatchGenerator from exo.worker.engines.mlx.generator.generate import ( PrefillCancelled, mlx_generate, warmup_inference, ) from exo.worker.engines.mlx.utils_mlx import ( apply_chat_template, mx_all_gather_tasks, mx_any, ) from exo.worker.runner.bootstrap import logger from .model_output_parsers import apply_all_parsers from .tool_parsers import ToolParser class Cancelled: pass class Finished: pass class GeneratorQueue[T]: def __init__(self): self._q = deque[T]() def push(self, t: T): self._q.append(t) def gen(self) -> Generator[T | None]: while True: if len(self._q) == 0: yield None else: yield self._q.popleft() class InferenceGenerator(ABC): _cancelled_tasks: set[TaskId] def should_cancel(self, task_id: TaskId) -> bool: return ( task_id in self._cancelled_tasks or CANCEL_ALL_TASKS in self._cancelled_tasks ) @abstractmethod def warmup(self) -> None: ... @abstractmethod def submit( self, task: TextGeneration, ) -> None: ... @abstractmethod def step( self, ) -> Iterable[ tuple[TaskId, ToolCallResponse | GenerationResponse | Cancelled | Finished] ]: ... @abstractmethod def close(self) -> None: ... EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL" EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM" EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT" def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None: """Check for debug prompt triggers in the input.""" from exo.worker.engines.mlx.utils_mlx import mlx_force_oom if len(task_params.input) == 0: return prompt = task_params.input[0].content if not prompt: return if EXO_RUNNER_MUST_FAIL in prompt: raise Exception("Artificial runner exception - for testing purposes only.") if EXO_RUNNER_MUST_OOM in prompt: mlx_force_oom() if EXO_RUNNER_MUST_TIMEOUT in prompt: time.sleep(100) @dataclass(eq=False) class SequentialGenerator(InferenceGenerator): model: Model tokenizer: TokenizerWrapper group: mx.distributed.Group | None kv_prefix_cache: KVPrefixCache | None tool_parser: ToolParser | None model_id: ModelId device_rank: int cancel_receiver: MpReceiver[TaskId] event_sender: MpSender[Event] check_for_cancel_every: int = 50 _cancelled_tasks: set[TaskId] = field(default_factory=set, init=False) _maybe_queue: list[TextGeneration] = field(default_factory=list, init=False) _maybe_cancel: list[TextGeneration] = field(default_factory=list, init=False) _all_tasks: dict[TaskId, TextGeneration] = field(default_factory=dict, init=False) _queue: deque[TextGeneration] = field(default_factory=deque, init=False) _active: ( tuple[ TextGeneration, # mlx generator that does work Generator[GenerationResponse], # queue that the 1st generator should push to and 3rd generator should pull from GeneratorQueue[GenerationResponse], # generator to get parsed outputs Generator[GenerationResponse | ToolCallResponse | None], ] | None ) = field(default=None, init=False) def warmup(self): self.check_for_cancel_every = warmup_inference( model=self.model, tokenizer=self.tokenizer, group=self.group, model_id=self.model_id, ) def submit( self, task: TextGeneration, ) -> None: self._cancelled_tasks.discard(CANCEL_ALL_TASKS) self._all_tasks[task.task_id] = task self._maybe_queue.append(task) def agree_on_tasks(self) -> None: """Agree between all ranks about the task ordering (some may have received in different order or not at all).""" agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group) self._queue.extend(task for task in self._maybe_queue if task in agreed) self._maybe_queue = [task for task in self._maybe_queue if task in different] def agree_on_cancellations(self) -> None: """Agree between all ranks about which tasks to cancel.""" has_cancel_all = False for task_id in self.cancel_receiver.collect(): if task_id == CANCEL_ALL_TASKS: has_cancel_all = True continue if task_id in self._all_tasks: self._maybe_cancel.append(self._all_tasks[task_id]) if mx_any(has_cancel_all, self.group): self._cancelled_tasks.add(CANCEL_ALL_TASKS) agreed, different = mx_all_gather_tasks(self._maybe_cancel, self.group) self._cancelled_tasks.update(task.task_id for task in agreed) self._maybe_cancel = list(different) def step( self, ) -> Iterable[ tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished] ]: if self._active is None: self.agree_on_tasks() if self._queue: self._start_next() else: return map(lambda task: (task, Cancelled()), self._cancelled_tasks) assert self._active is not None task, mlx_gen, queue, output_generator = self._active response = None try: queue.push(next(mlx_gen)) response = next(output_generator) except (StopIteration, PrefillCancelled): response = Finished() self._active = None if self._queue: self._start_next() except Exception as e: self._send_error(task, e) self._active = None raise return itertools.chain( [] if response is None else [(task.task_id, response)], map(lambda task: (task, Cancelled()), self._cancelled_tasks), ) def _start_next(self) -> None: task = self._queue.popleft() try: mlx_gen = self._build_generator(task) except Exception as e: self._send_error(task, e) raise queue = GeneratorQueue[GenerationResponse]() if task.task_params.bench: output_generator = queue.gen() else: output_generator = apply_all_parsers( queue.gen(), apply_chat_template(self.tokenizer, task.task_params), self.tool_parser, self.tokenizer, type(self.model), self.model_id, task.task_params.tools, ) self._active = (task, mlx_gen, queue, output_generator) def _send_error(self, task: TextGeneration, e: Exception) -> None: if self.device_rank == 0: self.event_sender.send( ChunkGenerated( command_id=task.command_id, chunk=ErrorChunk( model=self.model_id, finish_reason="error", error_message=str(e), ), ) ) def _build_generator(self, task: TextGeneration) -> Generator[GenerationResponse]: _check_for_debug_prompts(task.task_params) prompt = apply_chat_template(self.tokenizer, task.task_params) def on_prefill_progress(processed: int, total: int) -> None: if self.device_rank == 0: self.event_sender.send( ChunkGenerated( command_id=task.command_id, chunk=PrefillProgressChunk( model=self.model_id, processed_tokens=processed, total_tokens=total, ), ) ) def distributed_prompt_progress_callback() -> None: self.agree_on_cancellations() if self.should_cancel(task.task_id): raise PrefillCancelled() self.agree_on_tasks() tokens_since_cancel_check = self.check_for_cancel_every def on_generation_token() -> None: nonlocal tokens_since_cancel_check tokens_since_cancel_check += 1 if tokens_since_cancel_check >= self.check_for_cancel_every: tokens_since_cancel_check = 0 self.agree_on_cancellations() if self.should_cancel(task.task_id): raise PrefillCancelled() self.agree_on_tasks() return mlx_generate( model=self.model, tokenizer=self.tokenizer, task=task.task_params, prompt=prompt, kv_prefix_cache=self.kv_prefix_cache, on_prefill_progress=on_prefill_progress, distributed_prompt_progress_callback=distributed_prompt_progress_callback, on_generation_token=on_generation_token, group=self.group, ) def close(self) -> None: del self.model, self.tokenizer, self.group @dataclass(eq=False) class BatchGenerator(InferenceGenerator): model: Model tokenizer: TokenizerWrapper group: mx.distributed.Group | None kv_prefix_cache: KVPrefixCache | None tool_parser: ToolParser | None model_id: ModelId device_rank: int cancel_receiver: MpReceiver[TaskId] event_sender: MpSender[Event] check_for_cancel_every: int = 50 _cancelled_tasks: set[TaskId] = field(default_factory=set, init=False) _maybe_queue: list[TextGeneration] = field(default_factory=list, init=False) _maybe_cancel: list[TextGeneration] = field(default_factory=list, init=False) _all_tasks: dict[TaskId, TextGeneration] = field(default_factory=dict, init=False) _queue: deque[TextGeneration] = field(default_factory=deque, init=False) _mlx_gen: ExoBatchGenerator = field(init=False) _active_tasks: dict[ int, tuple[ TextGeneration, GeneratorQueue[GenerationResponse], Generator[GenerationResponse | ToolCallResponse | None], ], ] = field(default_factory=dict, init=False) def __post_init__(self) -> None: self._mlx_gen = ExoBatchGenerator( model=self.model, tokenizer=self.tokenizer, group=self.group, kv_prefix_cache=self.kv_prefix_cache, ) def warmup(self): self.check_for_cancel_every = warmup_inference( model=self.model, tokenizer=self.tokenizer, group=self.group, model_id=self.model_id, ) def submit( self, task: TextGeneration, ) -> None: self._cancelled_tasks.discard(CANCEL_ALL_TASKS) self._all_tasks[task.task_id] = task self._maybe_queue.append(task) def agree_on_tasks(self) -> None: """Agree between all ranks about the task ordering (some may have received in different order or not at all).""" agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group) self._queue.extend(task for task in self._maybe_queue if task in agreed) self._maybe_queue = [task for task in self._maybe_queue if task in different] def agree_on_cancellations(self) -> None: """Agree between all ranks about which tasks to cancel.""" has_cancel_all = False for task_id in self.cancel_receiver.collect(): if task_id == CANCEL_ALL_TASKS: has_cancel_all = True continue if task_id in self._all_tasks: self._maybe_cancel.append(self._all_tasks[task_id]) if mx_any(has_cancel_all, self.group): self._cancelled_tasks.add(CANCEL_ALL_TASKS) agreed, different = mx_all_gather_tasks(self._maybe_cancel, self.group) self._cancelled_tasks.update(task.task_id for task in agreed) self._maybe_cancel = list(different) def step( self, ) -> Iterable[ tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished] ]: if not self._queue: self.agree_on_tasks() # Submit any queued tasks to the engine while self._queue and len(self._active_tasks) < EXO_MAX_CONCURRENT_REQUESTS: task = self._queue.popleft() try: uid = self._start_task(task) except PrefillCancelled: continue except Exception as e: self._send_error(task, e) raise queue = GeneratorQueue[GenerationResponse]() if task.task_params.bench: output_generator = queue.gen() else: output_generator = apply_all_parsers( queue.gen(), apply_chat_template(self.tokenizer, task.task_params), self.tool_parser, self.tokenizer, type(self.model), self.model_id, task.task_params.tools, ) self._active_tasks[uid] = (task, queue, output_generator) if not self._mlx_gen.has_work: return self._apply_cancellations() results = self._mlx_gen.step() output: list[ tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished] ] = [] for uid, response in results: if uid not in self._active_tasks: # should we error here? logger.warning(f"{uid=} not found in active tasks") continue task, queue, output_generator = self._active_tasks[uid] queue.push(response) parsed = next(output_generator) if parsed is not None: output.append((task.task_id, parsed)) if response.finish_reason is not None: output.append((task.task_id, Finished())) del self._active_tasks[uid] return itertools.chain(output, self._apply_cancellations()) def _apply_cancellations( self, ) -> list[tuple[TaskId, Cancelled]]: if not self._cancelled_tasks: return [] cancel_all = CANCEL_ALL_TASKS in self._cancelled_tasks uids_to_cancel: list[int] = [] results: list[tuple[TaskId, Cancelled]] = [] for uid, (task, _, _) in list(self._active_tasks.items()): if task.task_id in self._cancelled_tasks or cancel_all: uids_to_cancel.append(uid) results.append((task.task_id, Cancelled())) del self._active_tasks[uid] if uids_to_cancel: self._mlx_gen.cancel(uids_to_cancel) already_cancelled = {tid for tid, _ in results} for tid in self._cancelled_tasks: if tid != CANCEL_ALL_TASKS and tid not in already_cancelled: results.append((tid, Cancelled())) self._cancelled_tasks.clear() return results def _send_error(self, task: TextGeneration, e: Exception) -> None: if self.device_rank == 0: self.event_sender.send( ChunkGenerated( command_id=task.command_id, chunk=ErrorChunk( model=self.model_id, finish_reason="error", error_message=str(e), ), ) ) def _start_task(self, task: TextGeneration) -> int: _check_for_debug_prompts(task.task_params) prompt = apply_chat_template(self.tokenizer, task.task_params) def on_prefill_progress(processed: int, total: int) -> None: if self.device_rank == 0: self.event_sender.send( ChunkGenerated( command_id=task.command_id, chunk=PrefillProgressChunk( model=self.model_id, processed_tokens=processed, total_tokens=total, ), ) ) def distributed_prompt_progress_callback() -> None: self.agree_on_cancellations() if self.should_cancel(task.task_id): raise PrefillCancelled() self.agree_on_tasks() tokens_since_cancel_check = self.check_for_cancel_every def on_generation_token() -> None: nonlocal tokens_since_cancel_check tokens_since_cancel_check += 1 if tokens_since_cancel_check >= self.check_for_cancel_every: tokens_since_cancel_check = 0 self.agree_on_cancellations() if self.should_cancel(task.task_id): self._cancelled_tasks.add(task.task_id) self.agree_on_tasks() return self._mlx_gen.submit( task_params=task.task_params, prompt=prompt, on_prefill_progress=on_prefill_progress, distributed_prompt_progress_callback=distributed_prompt_progress_callback, on_generation_token=on_generation_token, ) def close(self) -> None: self._mlx_gen.close() del self.model, self.tokenizer, self.group ================================================ FILE: src/exo/worker/runner/llm_inference/model_output_parsers.py ================================================ from collections.abc import Generator from functools import cache from typing import Any from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model from mlx_lm.models.gpt_oss import Model as GptOssModel from mlx_lm.tokenizer_utils import TokenizerWrapper from openai_harmony import ( # pyright: ignore[reportMissingTypeStubs] HarmonyEncodingName, HarmonyError, # pyright: ignore[reportUnknownVariableType] Role, StreamableParser, load_harmony_encoding, ) from exo.api.types import ToolCallItem from exo.shared.types.common import ModelId from exo.shared.types.mlx import Model from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse from exo.worker.engines.mlx.utils_mlx import ( detect_thinking_prompt_suffix, ) from exo.worker.runner.bootstrap import logger from exo.worker.runner.llm_inference.tool_parsers import ToolParser @cache def get_gpt_oss_encoding(): encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS) return encoding def apply_all_parsers( receiver: Generator[GenerationResponse | None], prompt: str, tool_parser: ToolParser | None, tokenizer: TokenizerWrapper, model_type: type[Model], model_id: ModelId, tools: list[dict[str, Any]] | None, ) -> Generator[GenerationResponse | ToolCallResponse | None]: mlx_generator = receiver if tokenizer.has_thinking: mlx_generator = parse_thinking_models( mlx_generator, tokenizer.think_start, tokenizer.think_end, starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer), ) if issubclass(model_type, GptOssModel): mlx_generator = parse_gpt_oss(mlx_generator) elif ( issubclass(model_type, DeepseekV32Model) and "deepseek" in model_id.normalize().lower() ): mlx_generator = parse_deepseek_v32(mlx_generator) elif tool_parser: mlx_generator = parse_tool_calls(mlx_generator, tool_parser, tools) return mlx_generator def parse_gpt_oss( responses: Generator[GenerationResponse | None], ) -> Generator[GenerationResponse | ToolCallResponse | None]: encoding = get_gpt_oss_encoding() stream = StreamableParser(encoding, role=Role.ASSISTANT) thinking = False current_tool_name: str | None = None tool_arg_parts: list[str] = [] for response in responses: if response is None: yield None continue try: stream.process(response.token) except HarmonyError: logger.error("Encountered critical Harmony Error, returning early") return delta = stream.last_content_delta ch = stream.current_channel recipient = stream.current_recipient # Debug: log every token with state logger.debug( f"parse_gpt_oss token={response.token} text={response.text!r} " f"recipient={recipient!r} ch={ch!r} delta={delta!r} " f"state={stream.state} current_tool={current_tool_name!r}" ) if recipient != current_tool_name: if current_tool_name is not None: prefix = "functions." if current_tool_name.startswith(prefix): current_tool_name = current_tool_name[len(prefix) :] logger.info( f"parse_gpt_oss yielding tool call: name={current_tool_name!r}" ) yield ToolCallResponse( tool_calls=[ ToolCallItem( name=current_tool_name, arguments="".join(tool_arg_parts).strip(), ) ], usage=response.usage, ) tool_arg_parts = [] current_tool_name = recipient # If inside a tool call, accumulate arguments if current_tool_name is not None: if delta: tool_arg_parts.append(delta) if response.finish_reason is not None: yield response.model_copy(update={"text": "".join(tool_arg_parts)}) tool_arg_parts = [] continue if ch == "analysis" and not thinking: thinking = True if ch != "analysis" and thinking: thinking = False if delta: yield response.model_copy(update={"text": delta, "is_thinking": thinking}) if response.finish_reason is not None: yield response def parse_deepseek_v32( responses: Generator[GenerationResponse | None], ) -> Generator[GenerationResponse | ToolCallResponse | None]: """Parse DeepSeek V3.2 DSML tool calls from the generation stream. Uses accumulated-text matching (not per-token marker checks) because DSML markers like <｜DSML｜function_calls> may span multiple tokens. Also handles ... blocks for thinking mode. """ from exo.worker.engines.mlx.dsml_encoding import ( THINKING_END, THINKING_START, TOOL_CALLS_END, TOOL_CALLS_START, parse_dsml_output, ) accumulated = "" in_tool_call = False thinking = False # Tokens buffered while we detect the start of a DSML block pending_buffer: list[GenerationResponse] = [] # Text accumulated during a tool call block tool_call_text = "" for response in responses: if response is None: yield None continue # ── Handle thinking tags ── if not thinking and THINKING_START in response.text: thinking = True # Yield any text before the tag before = response.text[: response.text.index(THINKING_START)] if before: yield response.model_copy(update={"text": before}) continue if thinking and THINKING_END in response.text: thinking = False # Yield any text after the tag after = response.text[ response.text.index(THINKING_END) + len(THINKING_END) : ] if after: yield response.model_copy(update={"text": after, "is_thinking": False}) continue if thinking: yield response.model_copy(update={"is_thinking": True}) continue # ── Handle tool call accumulation ── if in_tool_call: tool_call_text += response.text if TOOL_CALLS_END in tool_call_text: # Parse the accumulated DSML block parsed = parse_dsml_output(tool_call_text) if parsed is not None: logger.info(f"parsed DSML tool calls: {parsed}") yield ToolCallResponse( tool_calls=parsed, usage=response.usage, stats=response.stats, ) else: logger.warning( f"DSML tool call parsing failed for: {tool_call_text}" ) yield response.model_copy(update={"text": tool_call_text}) in_tool_call = False tool_call_text = "" continue # EOS reached before end marker — yield buffered text as-is if response.finish_reason is not None: logger.info("DSML tool call parsing interrupted by EOS") yield response.model_copy(update={"text": tool_call_text}) in_tool_call = False tool_call_text = "" continue # ── Detect start of tool call block ── accumulated += response.text if TOOL_CALLS_START in accumulated: # The start marker might be split across pending_buffer + current token start_idx = accumulated.index(TOOL_CALLS_START) # Yield any pending tokens that are purely before the marker pre_text = accumulated[:start_idx] if pre_text: # Flush pending buffer tokens that contributed text before the marker for buf_resp in pending_buffer: if pre_text: chunk = buf_resp.text if len(chunk) <= len(pre_text): yield buf_resp pre_text = pre_text[len(chunk) :] else: yield buf_resp.model_copy(update={"text": pre_text}) pre_text = "" pending_buffer = [] tool_call_text = accumulated[start_idx:] accumulated = "" # Check if the end marker is already present (entire tool call in one token) if TOOL_CALLS_END in tool_call_text: parsed = parse_dsml_output(tool_call_text) if parsed is not None: logger.info(f"parsed DSML tool calls: {parsed}") yield ToolCallResponse( tool_calls=parsed, usage=response.usage, stats=response.stats, ) else: logger.warning( f"DSML tool call parsing failed for: {tool_call_text}" ) yield response.model_copy(update={"text": tool_call_text}) tool_call_text = "" else: in_tool_call = True continue # Check if accumulated text might be the start of a DSML marker # Buffer tokens if we see a partial match at the end if _could_be_dsml_prefix(accumulated): pending_buffer.append(response) continue # No partial match — flush all pending tokens and the current one for buf_resp in pending_buffer: yield buf_resp pending_buffer = [] accumulated = "" yield response # Flush any remaining pending buffer at generator end for buf_resp in pending_buffer: yield buf_resp def _could_be_dsml_prefix(text: str) -> bool: """Check if the end of text could be the start of a DSML function_calls marker. We look for suffixes of text that are prefixes of the TOOL_CALLS_START pattern. This allows us to buffer tokens until we can determine if a tool call is starting. """ from exo.worker.engines.mlx.dsml_encoding import TOOL_CALLS_START # Only check the last portion of text that could overlap with the marker max_check = len(TOOL_CALLS_START) tail = text[-max_check:] if len(text) > max_check else text # Check if any suffix of tail is a prefix of TOOL_CALLS_START for i in range(len(tail)): suffix = tail[i:] if TOOL_CALLS_START.startswith(suffix): return True return False def parse_thinking_models( responses: Generator[GenerationResponse | None], think_start: str | None, think_end: str | None, starts_in_thinking: bool = True, ) -> Generator[GenerationResponse | None]: """Route thinking tokens via is_thinking flag. Swallows think tag tokens, sets is_thinking on all others. Always yields tokens with finish_reason to avoid hanging the chunk stream. """ is_thinking = starts_in_thinking for response in responses: if response is None: yield None continue if response.finish_reason is not None: yield response.model_copy(update={"is_thinking": False}) continue if response.text == think_start: is_thinking = True continue if response.text == think_end: is_thinking = False continue yield response.model_copy(update={"is_thinking": is_thinking}) def parse_tool_calls( responses: Generator[GenerationResponse | None], tool_parser: ToolParser, tools: list[dict[str, Any]] | None, ) -> Generator[GenerationResponse | ToolCallResponse | None]: in_tool_call = False tool_call_text_parts: list[str] = [] for response in responses: if response is None: yield None continue if not in_tool_call and response.text.startswith(tool_parser.start_parsing): in_tool_call = True if not in_tool_call: yield response continue tool_call_text_parts.append(response.text) if response.text.endswith(tool_parser.end_parsing): # parse the actual tool calls from the tool call text combined = "".join(tool_call_text_parts) parsed = tool_parser.parse(combined.strip(), tools=tools) logger.info(f"parsed {tool_call_text_parts=} into {parsed=}") in_tool_call = False tool_call_text_parts = [] if parsed is None: logger.warning(f"tool call parsing failed for text {combined}") yield response.model_copy(update={"text": combined}) continue yield ToolCallResponse( tool_calls=parsed, usage=response.usage, stats=response.stats ) continue if response.finish_reason is not None: logger.info( "tool call parsing interrupted, yield partial tool call as text" ) response = response.model_copy( update={ "text": "".join(tool_call_text_parts), "token": 0, } ) yield response ================================================ FILE: src/exo/worker/runner/llm_inference/runner.py ================================================ import os import time from dataclasses import dataclass from enum import Enum import mlx.core as mx from anyio import WouldBlock from mlx_lm.tokenizer_utils import TokenizerWrapper from exo.shared.models.model_cards import ModelTask from exo.shared.types.chunks import ( ErrorChunk, TokenChunk, ToolCallChunk, ) from exo.shared.types.common import CommandId, ModelId from exo.shared.types.events import ( ChunkGenerated, Event, RunnerStatusUpdated, TaskAcknowledged, TaskStatusUpdated, ) from exo.shared.types.mlx import Model from exo.shared.types.tasks import ( ConnectToGroup, LoadModel, Shutdown, StartWarmup, Task, TaskId, TaskStatus, TextGeneration, ) from exo.shared.types.worker.instances import BoundInstance from exo.shared.types.worker.runner_response import ( GenerationResponse, ToolCallResponse, ) from exo.shared.types.worker.runners import ( RunnerConnected, RunnerConnecting, RunnerFailed, RunnerIdle, RunnerLoaded, RunnerLoading, RunnerReady, RunnerRunning, RunnerShutdown, RunnerShuttingDown, RunnerStatus, RunnerWarmingUp, ) from exo.utils.channels import MpReceiver, MpSender from exo.worker.engines.mlx.cache import KVPrefixCache from exo.worker.engines.mlx.utils_mlx import ( initialize_mlx, load_mlx_items, ) from exo.worker.runner.bootstrap import logger from exo.worker.runner.llm_inference.batch_generator import ( BatchGenerator, InferenceGenerator, SequentialGenerator, ) from .batch_generator import Cancelled, Finished from .tool_parsers import make_mlx_parser class ExitCode(str, Enum): AllTasksComplete = "AllTasksComplete" Shutdown = "Shutdown" class Runner: def __init__( self, bound_instance: BoundInstance, event_sender: MpSender[Event], task_receiver: MpReceiver[Task], cancel_receiver: MpReceiver[TaskId], ): self.event_sender = event_sender self.task_receiver = task_receiver self.cancel_receiver = cancel_receiver self.bound_instance = bound_instance self.instance, self.runner_id, self.shard_metadata = ( self.bound_instance.instance, self.bound_instance.bound_runner_id, self.bound_instance.bound_shard, ) self.model_id = self.shard_metadata.model_card.model_id self.device_rank = self.shard_metadata.device_rank logger.info("hello from the runner") if getattr(self.shard_metadata, "immediate_exception", False): raise Exception("Fake exception - runner failed to spin up.") if timeout := getattr(self.shard_metadata, "should_timeout", 0): time.sleep(timeout) self.setup_start_time = time.time() self.generator: Builder | InferenceGenerator = Builder( self.model_id, self.event_sender, self.cancel_receiver ) self.seen: set[TaskId] = set() self.active_tasks: dict[ TaskId, TextGeneration, ] = {} logger.info("runner created") self.update_status(RunnerIdle()) def update_status(self, status: RunnerStatus): self.current_status = status self.event_sender.send( RunnerStatusUpdated( runner_id=self.runner_id, runner_status=self.current_status ) ) def send_task_status(self, task_id: TaskId, task_status: TaskStatus): self.event_sender.send( TaskStatusUpdated(task_id=task_id, task_status=task_status) ) def acknowledge_task(self, task: Task): self.event_sender.send(TaskAcknowledged(task_id=task.task_id)) def main(self): with self.task_receiver: for task in self.task_receiver: if task.task_id in self.seen: logger.warning("repeat task - potential error") continue self.seen.add(task.task_id) self.handle_first_task(task) if isinstance(self.current_status, RunnerShutdown): break def handle_first_task(self, task: Task): self.send_task_status(task.task_id, TaskStatus.Running) match task: case ConnectToGroup() if isinstance( self.current_status, (RunnerIdle, RunnerFailed) ): assert isinstance(self.generator, Builder) logger.info("runner connecting") self.update_status(RunnerConnecting()) self.acknowledge_task(task) self.generator.group = initialize_mlx(self.bound_instance) self.send_task_status(task.task_id, TaskStatus.Complete) self.update_status(RunnerConnected()) logger.info("runner connected") # we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to case LoadModel() if isinstance(self.generator, Builder) and ( ( isinstance(self.current_status, RunnerConnected) and self.generator.group is not None ) or ( isinstance(self.current_status, RunnerIdle) and self.generator.group is None ) ): total_layers = ( self.shard_metadata.end_layer - self.shard_metadata.start_layer ) logger.info("runner loading") self.update_status( RunnerLoading(layers_loaded=0, total_layers=total_layers) ) self.acknowledge_task(task) def on_model_load_timeout() -> None: self.update_status( RunnerFailed(error_message="Model loading timed out") ) time.sleep(0.5) def on_layer_loaded(layers_loaded: int, total: int) -> None: self.update_status( RunnerLoading(layers_loaded=layers_loaded, total_layers=total) ) assert ( ModelTask.TextGeneration in self.shard_metadata.model_card.tasks ), f"Incorrect model task(s): {self.shard_metadata.model_card.tasks}" self.generator.inference_model, self.generator.tokenizer = ( load_mlx_items( self.bound_instance, self.generator.group, on_timeout=on_model_load_timeout, on_layer_loaded=on_layer_loaded, ) ) self.generator = self.generator.build() self.send_task_status(task.task_id, TaskStatus.Complete) self.update_status(RunnerLoaded()) logger.info("runner loaded") case StartWarmup() if isinstance(self.current_status, RunnerLoaded): assert isinstance(self.generator, InferenceGenerator) logger.info("runner warming up") self.update_status(RunnerWarmingUp()) self.acknowledge_task(task) self.generator.warmup() logger.info( f"runner initialized in {time.time() - self.setup_start_time} seconds" ) self.send_task_status(task.task_id, TaskStatus.Complete) self.update_status(RunnerReady()) logger.info("runner ready") case TextGeneration() if isinstance(self.current_status, RunnerReady): return_code = self.handle_generation_tasks(starting_task=task) if return_code == ExitCode.Shutdown: return case Shutdown(): self.shutdown(task) return case _: raise ValueError( f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}" ) def shutdown(self, task: Task): logger.info("runner shutting down") self.update_status(RunnerShuttingDown()) self.acknowledge_task(task) if isinstance(self.generator, InferenceGenerator): self.generator.close() mx.clear_cache() import gc gc.collect() self.send_task_status(task.task_id, TaskStatus.Complete) self.update_status(RunnerShutdown()) def submit_text_generation(self, task: TextGeneration): assert isinstance(self.generator, InferenceGenerator) self.active_tasks[task.task_id] = task self.generator.submit(task) def handle_generation_tasks(self, starting_task: TextGeneration): assert isinstance(self.current_status, RunnerReady) assert isinstance(self.generator, InferenceGenerator) logger.info(f"received chat request: {starting_task}") self.update_status(RunnerRunning()) logger.info("runner running") self.acknowledge_task(starting_task) self.seen.add(starting_task.task_id) self.submit_text_generation(starting_task) while self.active_tasks: results = self.generator.step() finished: list[TaskId] = [] for task_id, result in results: match result: case Cancelled(): finished.append(task_id) case Finished(): self.send_task_status(task_id, TaskStatus.Complete) finished.append(task_id) case _: self.send_response( result, self.active_tasks[task_id].command_id ) for task_id in finished: self.active_tasks.pop(task_id, None) try: task = self.task_receiver.receive_nowait() if task.task_id in self.seen: logger.warning("repeat task - potential error") continue self.seen.add(task.task_id) match task: case TextGeneration(): self.acknowledge_task(task) self.submit_text_generation(task) case Shutdown(): self.shutdown(task) return ExitCode.Shutdown case _: raise ValueError( f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}" ) except WouldBlock: pass self.update_status(RunnerReady()) logger.info("runner ready") return ExitCode.AllTasksComplete def send_response( self, response: GenerationResponse | ToolCallResponse, command_id: CommandId ): match response: case GenerationResponse(): if self.device_rank == 0 and response.finish_reason == "error": self.event_sender.send( ChunkGenerated( command_id=command_id, chunk=ErrorChunk( error_message=response.text, model=self.model_id, ), ) ) elif self.device_rank == 0: assert response.finish_reason not in ( "error", "tool_calls", "function_call", ) self.event_sender.send( ChunkGenerated( command_id=command_id, chunk=TokenChunk( model=self.model_id, text=response.text, token_id=response.token, usage=response.usage, finish_reason=response.finish_reason, stats=response.stats, logprob=response.logprob, top_logprobs=response.top_logprobs, is_thinking=response.is_thinking, ), ) ) case ToolCallResponse(): if self.device_rank == 0: self.event_sender.send( ChunkGenerated( command_id=command_id, chunk=ToolCallChunk( tool_calls=response.tool_calls, model=self.model_id, usage=response.usage, stats=response.stats, ), ) ) @dataclass class Builder: model_id: ModelId event_sender: MpSender[Event] cancel_receiver: MpReceiver[TaskId] inference_model: Model | None = None tokenizer: TokenizerWrapper | None = None group: mx.distributed.Group | None = None def build( self, ) -> InferenceGenerator: assert self.model_id assert self.inference_model assert self.tokenizer tool_parser = None logger.info( f"model has_tool_calling={self.tokenizer.has_tool_calling} using tokens {self.tokenizer.tool_call_start}, {self.tokenizer.tool_call_end}" ) if ( self.tokenizer.tool_call_start and self.tokenizer.tool_call_end and self.tokenizer.tool_parser # type: ignore ): tool_parser = make_mlx_parser( self.tokenizer.tool_call_start, self.tokenizer.tool_call_end, self.tokenizer.tool_parser, # type: ignore ) kv_prefix_cache = KVPrefixCache(self.group) device_rank = 0 if self.group is None else self.group.rank() if os.environ.get("EXO_NO_BATCH"): logger.info("using SequentialGenerator (batching disabled)") return SequentialGenerator( model=self.inference_model, tokenizer=self.tokenizer, group=self.group, tool_parser=tool_parser, kv_prefix_cache=kv_prefix_cache, model_id=self.model_id, device_rank=device_rank, cancel_receiver=self.cancel_receiver, event_sender=self.event_sender, ) logger.info("using BatchGenerator") return BatchGenerator( model=self.inference_model, tokenizer=self.tokenizer, group=self.group, tool_parser=tool_parser, kv_prefix_cache=kv_prefix_cache, model_id=self.model_id, device_rank=device_rank, cancel_receiver=self.cancel_receiver, event_sender=self.event_sender, ) ================================================ FILE: src/exo/worker/runner/llm_inference/tool_parsers.py ================================================ import json import math from dataclasses import dataclass from typing import Any, Callable from exo.api.types import ToolCallItem @dataclass class ToolParser: start_parsing: str end_parsing: str _inner_parser: Callable[[str], list[ToolCallItem] | None] def parse( self, text: str, tools: list[dict[str, Any]] | None ) -> list[ToolCallItem] | None: parsed = self._inner_parser(text) if parsed is None: return None if tools is not None: parsed = _coerce_tool_calls_to_schema(parsed, tools) return parsed def _json_type_matches(value: Any, expected_type: str) -> bool: # pyright: ignore[reportAny] if expected_type == "object": return isinstance(value, dict) if expected_type == "array": return isinstance(value, list) if expected_type == "string": return isinstance(value, str) if expected_type == "integer": return isinstance(value, int) and not isinstance(value, bool) if expected_type == "number": return (isinstance(value, int) and not isinstance(value, bool)) or isinstance( value, float ) if expected_type == "boolean": return isinstance(value, bool) if expected_type == "null": return value is None return False def _coerce_tool_arg_with_schema(value: Any, schema: dict[str, Any]) -> Any: # pyright: ignore[reportAny] schema_type = schema.get("type") if isinstance(schema_type, list): for candidate in schema_type: # pyright: ignore[reportUnknownVariableType] if not isinstance(candidate, str): continue if candidate == "null" and value is None: return None candidate_schema = {**schema, "type": candidate} coerced = _coerce_tool_arg_with_schema(value, candidate_schema) # pyright: ignore[reportAny] if _json_type_matches(coerced, candidate): return coerced # pyright: ignore[reportAny] return value # pyright: ignore[reportAny] if not isinstance(schema_type, str): return value # pyright: ignore[reportAny] if schema_type == "object": parsed = value # pyright: ignore[reportAny] if isinstance(parsed, str): try: parsed = json.loads(parsed) # pyright: ignore[reportAny] except Exception: return value # pyright: ignore[reportAny] if not isinstance(parsed, dict): return value # pyright: ignore[reportAny] properties = schema.get("properties") if not isinstance(properties, dict): return parsed # pyright: ignore[reportUnknownVariableType] return { key: ( _coerce_tool_arg_with_schema(prop_value, prop_schema) # pyright: ignore[reportUnknownArgumentType] if isinstance(prop_schema, dict) else prop_value ) for key, prop_value in parsed.items() # pyright: ignore[reportUnknownVariableType] for prop_schema in [properties.get(key)] # type: ignore } if schema_type == "array": parsed = value # pyright: ignore[reportAny] if isinstance(parsed, str): try: parsed = json.loads(parsed) # pyright: ignore[reportAny] except Exception: return value # pyright: ignore[reportAny] if not isinstance(parsed, list): return value # pyright: ignore[reportAny] item_schema = schema.get("items") if not isinstance(item_schema, dict): return parsed # pyright: ignore[reportUnknownVariableType] return [_coerce_tool_arg_with_schema(item, item_schema) for item in parsed] # type: ignore if schema_type == "integer": if isinstance(value, bool): return value if isinstance(value, int): return value if isinstance(value, float) and value.is_integer(): return int(value) if isinstance(value, str): try: return int(value.strip()) except ValueError: return value return value if schema_type == "number": if isinstance(value, bool): return value if isinstance(value, (int, float)): return value if isinstance(value, str): try: num = float(value.strip()) if math.isfinite(num): return num except ValueError: return value return value if schema_type == "boolean": if isinstance(value, bool): return value if isinstance(value, str): lowered = value.strip().lower() if lowered == "true": return True if lowered == "false": return False return value return value # pyright: ignore[reportAny] def _coerce_tool_calls_to_schema( tool_calls: list[ToolCallItem], tools: list[dict[str, Any]] ) -> list[ToolCallItem]: schema_by_name: dict[str, dict[str, Any]] = {} for tool in tools: function = tool.get("function") if not isinstance(function, dict): continue name = function.get("name") # type: ignore parameters = function.get("parameters") # type: ignore if isinstance(name, str) and isinstance(parameters, dict): schema_by_name[name] = parameters if not schema_by_name: return tool_calls coerced_calls: list[ToolCallItem] = [] for tool_call in tool_calls: schema = schema_by_name.get(tool_call.name) if schema is None: coerced_calls.append(tool_call) continue try: parsed_args = json.loads(tool_call.arguments) # pyright: ignore[reportAny] except Exception: coerced_calls.append(tool_call) continue if not isinstance(parsed_args, dict): coerced_calls.append(tool_call) continue coerced_args = _coerce_tool_arg_with_schema(parsed_args, schema) # pyright: ignore[reportAny] if not isinstance(coerced_args, dict): coerced_calls.append(tool_call) continue coerced_calls.append( tool_call.model_copy(update={"arguments": json.dumps(coerced_args)}) ) return coerced_calls def make_mlx_parser( tool_call_start: str, tool_call_end: str, tool_parser: Callable[[str], dict[str, Any] | list[dict[str, Any]]], ) -> ToolParser: def parse_tool_calls(text: str) -> list[ToolCallItem] | None: try: text = text.removeprefix(tool_call_start) text = text.removesuffix(tool_call_end) parsed = tool_parser(text) if isinstance(parsed, list): return [ToolCallItem.model_validate(_flatten(p)) for p in parsed] else: return [ToolCallItem.model_validate(_flatten(parsed))] except Exception: return None return ToolParser( start_parsing=tool_call_start, end_parsing=tool_call_end, _inner_parser=parse_tool_calls, ) # TODO / example code: def _parse_json_calls(text: str) -> list[ToolCallItem] | None: try: text = text.removeprefix("") text = text.removesuffix("") top_level = { k: json.dumps(v) if isinstance(v, (dict, list)) else v for k, v in json.loads(text).items() # pyright: ignore[reportAny] } return [ToolCallItem.model_validate(top_level)] except Exception: return None def _flatten(p: dict[str, Any]) -> dict[str, str]: return { k: json.dumps(v) if isinstance(v, (dict, list)) else str(v) # pyright: ignore[reportAny] for k, v in p.items() # pyright: ignore[reportAny] } def make_json_parser() -> ToolParser: return ToolParser( start_parsing="", end_parsing="", _inner_parser=_parse_json_calls, ) def infer_tool_parser(chat_template: str) -> ToolParser | None: """Attempt to auto-infer a tool parser from the chat template.""" if "" in chat_template and "tool_call.name" in chat_template: return make_json_parser() return None ================================================ FILE: src/exo/worker/runner/runner_supervisor.py ================================================ import contextlib import multiprocessing as mp import signal from dataclasses import dataclass, field from typing import Self import anyio from anyio import ( BrokenResourceError, ClosedResourceError, to_thread, ) from loguru import logger from exo.shared.types.chunks import ErrorChunk from exo.shared.types.events import ( ChunkGenerated, Event, RunnerStatusUpdated, TaskAcknowledged, TaskStatusUpdated, ) from exo.shared.types.tasks import ( CANCEL_ALL_TASKS, ImageEdits, ImageGeneration, Task, TaskId, TaskStatus, TextGeneration, ) from exo.shared.types.worker.instances import BoundInstance from exo.shared.types.worker.runners import ( RunnerConnecting, RunnerFailed, RunnerIdle, RunnerLoading, RunnerRunning, RunnerShuttingDown, RunnerStatus, RunnerWarmingUp, ) from exo.shared.types.worker.shards import ShardMetadata from exo.utils.channels import MpReceiver, MpSender, Sender, mp_channel from exo.utils.task_group import TaskGroup from exo.worker.runner.bootstrap import entrypoint PREFILL_TIMEOUT_SECONDS = 60 DECODE_TIMEOUT_SECONDS = 5 @dataclass(eq=False) class RunnerSupervisor: shard_metadata: ShardMetadata bound_instance: BoundInstance runner_process: mp.Process initialize_timeout: float _ev_recv: MpReceiver[Event] _task_sender: MpSender[Task] _event_sender: Sender[Event] _cancel_sender: MpSender[TaskId] _tg: TaskGroup = field(default_factory=TaskGroup, init=False) status: RunnerStatus = field(default_factory=RunnerIdle, init=False) pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False) in_progress: dict[TaskId, Task] = field(default_factory=dict, init=False) completed: set[TaskId] = field(default_factory=set, init=False) cancelled: set[TaskId] = field(default_factory=set, init=False) _cancel_watch_runner: anyio.CancelScope = field( default_factory=anyio.CancelScope, init=False ) @classmethod def create( cls, *, bound_instance: BoundInstance, event_sender: Sender[Event], initialize_timeout: float = 400, ) -> Self: ev_send, ev_recv = mp_channel[Event]() task_sender, task_recv = mp_channel[Task]() cancel_sender, cancel_recv = mp_channel[TaskId]() runner_process = mp.Process( target=entrypoint, args=( bound_instance, ev_send, task_recv, cancel_recv, logger, ), daemon=True, ) shard_metadata = bound_instance.bound_shard self = cls( bound_instance=bound_instance, shard_metadata=shard_metadata, runner_process=runner_process, initialize_timeout=initialize_timeout, _ev_recv=ev_recv, _task_sender=task_sender, _cancel_sender=cancel_sender, _event_sender=event_sender, ) return self async def run(self): self.runner_process.start() async with self._tg as tg: tg.start_soon(self._watch_runner) tg.start_soon(self._forward_events) def shutdown(self): logger.info("Runner supervisor shutting down") self._tg.cancel_tasks() if not self._cancel_watch_runner.cancel_called: self._cancel_watch_runner.cancel() with contextlib.suppress(ClosedResourceError): self._ev_recv.close() with contextlib.suppress(ClosedResourceError): self._task_sender.close() with contextlib.suppress(ClosedResourceError): self._event_sender.close() with contextlib.suppress(ClosedResourceError): self._cancel_sender.send(CANCEL_ALL_TASKS) with contextlib.suppress(ClosedResourceError): self._cancel_sender.close() self.runner_process.join(5) if not self.runner_process.is_alive(): logger.info("Runner process succesfully terminated") return # This is overkill but it's not technically bad, just unnecessary. logger.warning("Runner process didn't shutdown succesfully, terminating") self.runner_process.terminate() self.runner_process.join(1) if not self.runner_process.is_alive(): return logger.critical("Runner process didn't respond to SIGTERM, killing") self.runner_process.kill() async def start_task(self, task: Task): if task.task_id in self.pending: logger.warning( f"Skipping invalid task {task} as it has already been submitted" ) return if task.task_id in self.completed: logger.warning( f"Skipping invalid task {task} as it has already been completed" ) return logger.info(f"Starting task {task}") event = anyio.Event() self.pending[task.task_id] = event self.in_progress[task.task_id] = task try: await self._task_sender.send_async(task) except ClosedResourceError: self.in_progress.pop(task.task_id, None) logger.warning(f"Task {task} dropped, runner closed communication.") return await event.wait() async def cancel_task(self, task_id: TaskId): if task_id in self.completed: logger.info(f"Unable to cancel {task_id} as it has been completed") self.cancelled.add(task_id) return self.cancelled.add(task_id) with anyio.move_on_after(0.5) as scope: try: await self._cancel_sender.send_async(task_id) except ClosedResourceError: # typically occurs when trying to shut down a failed instance logger.warning( f"Cancelling task {task_id} failed, runner closed communication" ) if scope.cancel_called: logger.error("RunnerSupervisor cancel pipe blocked") await self._check_runner(TimeoutError("cancel pipe blocked")) async def _forward_events(self): try: with self._ev_recv as events: async for event in events: if isinstance(event, RunnerStatusUpdated): self.status = event.runner_status if isinstance(event, TaskAcknowledged): self.pending.pop(event.task_id).set() continue if ( isinstance(event, TaskStatusUpdated) and event.task_status == TaskStatus.Complete ): # If a task has just been completed, we should be working on it. assert isinstance( self.status, ( RunnerRunning, RunnerWarmingUp, RunnerLoading, RunnerConnecting, RunnerShuttingDown, ), ) self.in_progress.pop(event.task_id, None) self.completed.add(event.task_id) await self._event_sender.send(event) except (ClosedResourceError, BrokenResourceError) as e: await self._check_runner(e) finally: for tid in self.pending: self.pending[tid].set() def __del__(self) -> None: if self.runner_process.is_alive(): logger.critical("RunnerSupervisor was not stopped cleanly.") with contextlib.suppress(ValueError): self.runner_process.kill() async def _watch_runner(self) -> None: with self._cancel_watch_runner: while True: await anyio.sleep(5) if not self.runner_process.is_alive(): await self._check_runner(RuntimeError("Runner found to be dead")) async def _check_runner(self, e: Exception) -> None: if not self._cancel_watch_runner.cancel_called: self._cancel_watch_runner.cancel() logger.info("Checking runner's status") if self.runner_process.is_alive(): logger.info("Runner was found to be alive, attempting to join process") await to_thread.run_sync(self.runner_process.join, 5) rc = self.runner_process.exitcode logger.info(f"Runner exited with exit code {rc}") if rc == 0: return if isinstance(rc, int) and rc < 0: sig = -rc try: cause = f"signal={sig} ({signal.strsignal(sig)})" except Exception: cause = f"signal={sig}" else: cause = f"exitcode={rc}" logger.opt(exception=e).error(f"Runner terminated with {cause}") for task in self.in_progress.values(): if isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)): with anyio.CancelScope(shield=True): await self._event_sender.send( ChunkGenerated( command_id=task.command_id, chunk=ErrorChunk( model=self.shard_metadata.model_card.model_id, error_message=( "Runner shutdown before completing command " f"({cause})" ), ), ) ) try: self.status = RunnerFailed(error_message=f"Terminated ({cause})") with anyio.CancelScope(shield=True): await self._event_sender.send( RunnerStatusUpdated( runner_id=self.bound_instance.bound_runner_id, runner_status=RunnerFailed( error_message=f"Terminated ({cause})" ), ) ) except (ClosedResourceError, BrokenResourceError): logger.warning( "Event sender already closed, unable to report runner failure" ) self.shutdown() ================================================ FILE: src/exo/worker/tests/TODO.tests ================================================ Unit Tests 1. Test worker plans as expected - State transitions are correct - Unexpected states throw 2. Test runner - Stays loaded - Unloads under end condition - Accepts tasks - Returns ChunkGenerated events 3. Test mlx engine - Autoparallel on n of the same nodes returns tensors with 1/n size - mx.barrier forces computation - Distributed init returns expected configuration - initialize_mlx sets wired limit - shard_and_load returns expected model - Quantization returns quantized layers 4. Download - hits the correct endpoint - normalizes tags correctly - updates download progress 5. Serialization/Deserialization of tagged models Integration tests: 1. Test model inference is "sensible" (per-configuration) - Non-empty response - Sensible inference speed - Answers are non-gibberish for many seeds (What is the capital of France? -> "Paris" in answer.) - Answer is the same for particular seed 2. Test that node count does not affect inference result (per-configuration) - Llama on 1 node, and on 2 nodes returns the same result, given temperature 0 and set seed. - Do for all configurations (Ring/Jaccl, Pipeline/Tensor) 3. Test supervisor catches exceptions gracefully - Timeouts - OOM - MLX error 4. distributed init memory requirements are as expected 5. MLX - KVCache size is same length as prompt tokens - Prefix cache (once implemented) 6. Spin up creates a runner or goes to failed status Regression tests: 1. Per-configuration baseline performance - no 20% drop in performance (device, node count, model, strategy, backend) ================================================ FILE: src/exo/worker/tests/__init__.py ================================================ ================================================ FILE: src/exo/worker/tests/constants.py ================================================ from typing import Final from exo.shared.models.model_cards import ModelId from exo.shared.types.common import CommandId, NodeId from exo.shared.types.tasks import TaskId from exo.shared.types.worker.instances import InstanceId, RunnerId MASTER_NODE_ID = NodeId("ffffffff-aaaa-4aaa-8aaa-aaaaaaaaaaaa") NODE_A: Final[NodeId] = NodeId("aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa") NODE_B: Final[NodeId] = NodeId("bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb") NODE_C: Final[NodeId] = NodeId("cccccccc-cccc-4ccc-8ccc-cccccccccccc") RUNNER_1_ID: Final[RunnerId] = RunnerId("11111111-1111-4111-8111-111111111111") RUNNER_2_ID: Final[RunnerId] = RunnerId("33333333-3333-4333-8333-333333333333") RUNNER_3_ID: Final[RunnerId] = RunnerId("Runner3") INSTANCE_1_ID: Final[InstanceId] = InstanceId("22222222-2222-4222-8222-222222222222") INSTANCE_2_ID: Final[InstanceId] = InstanceId("44444444-4444-4444-8444-444444444444") MODEL_A_ID: Final[ModelId] = ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit") MODEL_B_ID: Final[ModelId] = ModelId("mlx-community/TinyLlama-1.1B-Chat-v1.0") TASK_1_ID: Final[TaskId] = TaskId("55555555-5555-4555-8555-555555555555") TASK_2_ID: Final[TaskId] = TaskId("66666666-6666-4666-8666-666666666666") COMMAND_1_ID: Final[CommandId] = CommandId("77777777-7777-4777-8777-777777777777") COMMAND_2_ID: Final[CommandId] = CommandId("88888888-8888-4888-8888-888888888888") SHUTDOWN_TASK_ID = TaskId("shutdown") CHAT_COMPLETION_TASK_ID = TaskId("chat-completion") INITIALIZATION_TASK_ID = TaskId("initialisation") LOAD_TASK_ID = TaskId("load") WARMUP_TASK_ID = TaskId("warmup") ================================================ FILE: src/exo/worker/tests/unittests/__init__.py ================================================ ================================================ FILE: src/exo/worker/tests/unittests/conftest.py ================================================ from dataclasses import dataclass, field from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask from exo.shared.types.common import NodeId from exo.shared.types.memory import Memory from exo.shared.types.tasks import BaseTask, TaskId from exo.shared.types.worker.instances import ( BoundInstance, Instance, InstanceId, MlxRingInstance, ) from exo.shared.types.worker.runners import RunnerId, RunnerStatus, ShardAssignments from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata # Runner supervisor without multiprocessing logic. @dataclass(frozen=True) class FakeRunnerSupervisor: bound_instance: BoundInstance status: RunnerStatus completed: set[TaskId] = field(default_factory=set) in_progress: set[TaskId] = field(default_factory=set) pending: dict[TaskId, object] = field(default_factory=dict) class OtherTask(BaseTask): pass # TODO: Is this actually better than using Mock/Fake dataclasses? # e.g. commit d01cd292344df15759070966826a6c027945792b def get_pipeline_shard_metadata( model_id: ModelId, device_rank: int, world_size: int = 1 ) -> ShardMetadata: return PipelineShardMetadata( model_card=ModelCard( model_id=model_id, storage_size=Memory.from_mb(100000), n_layers=32, hidden_size=2048, supports_tensor=False, tasks=[ModelTask.TextGeneration], ), device_rank=device_rank, world_size=world_size, start_layer=0, end_layer=32, n_layers=32, ) def get_shard_assignments( model_id: ModelId, node_to_runner: dict[NodeId, RunnerId], runner_to_shard: dict[RunnerId, ShardMetadata], ) -> ShardAssignments: return ShardAssignments( model_id=model_id, node_to_runner=node_to_runner, runner_to_shard=runner_to_shard, ) def get_mlx_ring_instance( instance_id: InstanceId, model_id: ModelId, node_to_runner: dict[NodeId, RunnerId], runner_to_shard: dict[RunnerId, ShardMetadata], ) -> Instance: return MlxRingInstance( instance_id=instance_id, shard_assignments=get_shard_assignments( model_id, node_to_runner, runner_to_shard ), hosts_by_node={}, ephemeral_port=50000, ) def get_bound_mlx_ring_instance( instance_id: InstanceId, model_id: ModelId, runner_id: RunnerId, node_id: NodeId ) -> BoundInstance: shard = get_pipeline_shard_metadata(model_id=model_id, device_rank=0, world_size=2) other_shard = get_pipeline_shard_metadata( model_id=model_id, device_rank=1, world_size=2 ) instance = get_mlx_ring_instance( instance_id=instance_id, model_id=model_id, node_to_runner={ node_id: runner_id, NodeId("other_node"): RunnerId("other_runner"), }, runner_to_shard={runner_id: shard, RunnerId("other_runner"): other_shard}, ) return BoundInstance( instance=instance, bound_runner_id=runner_id, bound_node_id=node_id ) ================================================ FILE: src/exo/worker/tests/unittests/test_download/__init__.py ================================================ ================================================ FILE: src/exo/worker/tests/unittests/test_mlx/__init__.py ================================================ ================================================ FILE: src/exo/worker/tests/unittests/test_mlx/conftest.py ================================================ # type: ignore import json import os import tempfile import traceback from dataclasses import dataclass from pathlib import Path from typing import Any, cast import mlx.core as mx import mlx.nn as nn from exo.shared.constants import EXO_MODELS_DIR from exo.shared.models.model_cards import ModelCard, ModelTask from exo.shared.types.common import ModelId from exo.shared.types.memory import Memory from exo.shared.types.mlx import Model from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata from exo.worker.engines.mlx.generator.generate import mlx_generate from exo.worker.engines.mlx.utils_mlx import apply_chat_template, shard_and_load class MockLayer(nn.Module): def __init__(self) -> None: super().__init__() self.custom_attr = "test_value" self.use_sliding = True def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array: return x * 2 @dataclass(frozen=True) class PipelineTestConfig: model_path: Path total_layers: int base_port: int max_tokens: int def create_hostfile(world_size: int, base_port: int) -> tuple[str, list[str]]: hosts = [f"127.0.0.1:{base_port + i}" for i in range(world_size)] with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: json.dump(hosts, f) hostfile_path = f.name return hostfile_path, hosts # Use GPT OSS 20b to test as it is a model with a lot of strange behaviour DEFAULT_GPT_OSS_CONFIG = PipelineTestConfig( model_path=EXO_MODELS_DIR / "mlx-community--gpt-oss-20b-MXFP4-Q8", total_layers=24, base_port=29600, max_tokens=200, ) DEFAULT_GPT_OSS_MODEL_ID = "mlx-community/gpt-oss-20b-MXFP4-Q8" def run_gpt_oss_pipeline_device( rank: int, world_size: int, hostfile_path: str, layer_splits: list[tuple[int, int]], prompt_tokens: int, prefill_step_size: int, result_queue: Any, # pyright: ignore[reportAny] max_tokens: int = 200, ) -> None: os.environ["MLX_HOSTFILE"] = hostfile_path os.environ["MLX_RANK"] = str(rank) try: group = mx.distributed.init(backend="ring", strict=True) start_layer, end_layer = layer_splits[rank] shard_meta = PipelineShardMetadata( model_card=ModelCard( model_id=ModelId(DEFAULT_GPT_OSS_MODEL_ID), storage_size=Memory.from_gb(12), n_layers=24, hidden_size=2880, supports_tensor=False, tasks=[ModelTask.TextGeneration], ), device_rank=rank, world_size=world_size, start_layer=start_layer, end_layer=end_layer, n_layers=24, ) model, tokenizer = shard_and_load( shard_meta, group, on_timeout=None, on_layer_loaded=None ) model = cast(Model, model) # Generate a prompt of exact token length base_text = "The quick brown fox jumps over the lazy dog. " base_tokens = tokenizer.encode(base_text) base_len = len(base_tokens) # Build prompt with approximate target length repeats = (prompt_tokens // base_len) + 2 long_text = base_text * repeats tokens = tokenizer.encode(long_text) # Truncate to exact target length tokens = tokens[:prompt_tokens] prompt_text = tokenizer.decode(tokens) task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content=prompt_text)], max_output_tokens=max_tokens, ) prompt = apply_chat_template(tokenizer, task) generated_text = "" for response in mlx_generate( model=model, tokenizer=tokenizer, task=task, prompt=prompt, kv_prefix_cache=None, group=group, ): generated_text += response.text if response.finish_reason is not None: break result_queue.put((rank, True, generated_text)) # pyright: ignore[reportAny] except Exception as e: result_queue.put((rank, False, f"{e}\n{traceback.format_exc()}")) # pyright: ignore[reportAny] def run_gpt_oss_tensor_parallel_device( rank: int, world_size: int, hostfile_path: str, prompt_tokens: int, prefill_step_size: int, result_queue: Any, # pyright: ignore[reportAny] max_tokens: int = 10, ) -> None: os.environ["MLX_HOSTFILE"] = hostfile_path os.environ["MLX_RANK"] = str(rank) try: group = mx.distributed.init(backend="ring", strict=True) # For tensor parallelism, all devices run all layers shard_meta = TensorShardMetadata( model_card=ModelCard( model_id=ModelId(DEFAULT_GPT_OSS_MODEL_ID), storage_size=Memory.from_gb(12), n_layers=24, hidden_size=2880, supports_tensor=True, tasks=[ModelTask.TextGeneration], ), device_rank=rank, world_size=world_size, start_layer=0, end_layer=24, n_layers=24, ) model, tokenizer = shard_and_load( shard_meta, group, on_timeout=None, on_layer_loaded=None ) model = cast(Model, model) base_text = "The quick brown fox jumps over the lazy dog. " base_tokens = tokenizer.encode(base_text) base_len = len(base_tokens) repeats = (prompt_tokens // base_len) + 2 long_text = base_text * repeats tokens = tokenizer.encode(long_text) tokens = tokens[:prompt_tokens] prompt_text = tokenizer.decode(tokens) task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content=prompt_text)], max_output_tokens=max_tokens, ) prompt = apply_chat_template(tokenizer, task) generated_text = "" for response in mlx_generate( model=model, tokenizer=tokenizer, task=task, prompt=prompt, kv_prefix_cache=None, group=group, ): generated_text += response.text if response.finish_reason is not None: break result_queue.put((rank, True, generated_text)) # pyright: ignore[reportAny] except Exception as e: result_queue.put((rank, False, f"{e}\n{traceback.format_exc()}")) # pyright: ignore[reportAny] ================================================ FILE: src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py ================================================ import json import multiprocessing as mp import os import tempfile from typing import Any import mlx.core as mx import mlx.nn as mlx_nn import pytest from exo.worker.engines.mlx.auto_parallel import ( CustomMlxLayer, PipelineFirstLayer, PipelineLastLayer, patch_pipeline_model, ) from exo.worker.tests.unittests.test_mlx.conftest import MockLayer def run_pipeline_device( rank: int, world_size: int, hostfile_path: str, result_queue: Any, # pyright: ignore[reportAny] ) -> None: import os os.environ["MLX_HOSTFILE"] = hostfile_path os.environ["MLX_RANK"] = str(rank) class MockLayerInner(mlx_nn.Module): def __init__(self) -> None: super().__init__() self.custom_attr = "test_value" def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array: return x * 2 class MockModel(mlx_nn.Module): def __init__(self, layers: list[mlx_nn.Module]) -> None: super().__init__() self.layers = layers def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array: for layer in self.layers: x = layer(x, *args, **kwargs) # pyright: ignore[reportUnknownVariableType] return x # pyright: ignore[reportUnknownVariableType] try: group = mx.distributed.init(backend="ring", strict=True) mock = MockLayerInner() first = PipelineFirstLayer(mock, r=rank, group=group) composed = PipelineLastLayer(first, r=rank, s=world_size, group=group) # Wrap in a mock model, then wrap in PipelineParallelModel for all_gather inner_model = MockModel([composed]) model = patch_pipeline_model(inner_model, group) x = mx.ones((1, 4)) result = model(x) mx.eval(result) success = result.shape == x.shape result_queue.put((rank, success, result)) # pyright: ignore[reportAny] except Exception as e: result_queue.put((rank, False, str(e))) # pyright: ignore[reportAny] def test_single_wrapper_delegates_attributes() -> None: mock = MockLayer() wrapped = CustomMlxLayer(mock) assert wrapped.custom_attr == "test_value" # type: ignore[attr-defined] assert wrapped.use_sliding is True # type: ignore[attr-defined] def test_composed_wrappers_delegate_attributes() -> None: mock = MockLayer() group = mx.distributed.init() first = PipelineFirstLayer(mock, r=0, group=group) composed = PipelineLastLayer(first, r=0, s=1, group=group) assert composed.custom_attr == "test_value" # type: ignore[attr-defined] assert composed.use_sliding is True # type: ignore[attr-defined] def test_missing_attribute_raises() -> None: mock = MockLayer() wrapped = CustomMlxLayer(mock) with pytest.raises(AttributeError): _ = wrapped.nonexistent_attr # type: ignore[attr-defined] def test_composed_call_works() -> None: ctx = mp.get_context("spawn") world_size = 2 base_port = 29500 hosts = [f"127.0.0.1:{base_port + i}" for i in range(world_size)] with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: json.dump(hosts, f) hostfile_path = f.name try: result_queue: Any = ctx.Queue() processes: list[Any] = [] for rank in range(world_size): p = ctx.Process( target=run_pipeline_device, args=(rank, world_size, hostfile_path, result_queue), ) p.start() processes.append(p) for p in processes: # pyright: ignore[reportAny] p.join(timeout=10) # pyright: ignore[reportAny] results: dict[int, Any] = {} errors: dict[int, str] = {} while not result_queue.empty(): # pyright: ignore[reportAny] rank, success, value = result_queue.get() # pyright: ignore[reportAny] if success: results[rank] = value else: errors[rank] = value assert len(results) == world_size, ( f"Expected {world_size} results, got {len(results)}. Errors: {errors}" ) for rank in range(world_size): assert rank in results, ( f"Device {rank} failed: {errors.get(rank, 'unknown')}" ) result_array = results[rank] # Both devices see the final result (4.0) after all_gather assert (result_array == 4.0).all(), ( f"Device {rank}: expected 4.0, got {result_array}" ) finally: os.unlink(hostfile_path) ================================================ FILE: src/exo/worker/tests/unittests/test_mlx/test_batch_vs_generate.py ================================================ import copy import gc import json import shutil import tempfile from pathlib import Path from typing import Any, cast import mlx.core as mx import pytest from mlx_lm.tokenizer_utils import TokenizerWrapper from exo.shared.types.common import ModelId from exo.shared.types.mlx import KVCacheType, Model from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams from exo.worker.engines.mlx.cache import CacheSnapshot, KVPrefixCache, cache_length from exo.worker.engines.mlx.generator.batch_generate import ExoBatchGenerator from exo.worker.engines.mlx.generator.generate import mlx_generate from exo.worker.engines.mlx.utils_mlx import ( apply_chat_template, load_tokenizer_for_model_id, ) from .test_prefix_cache_architectures import ( ARCHITECTURES, ArchSpec, _arch_available, # pyright: ignore[reportPrivateUsage] _build_model, # pyright: ignore[reportPrivateUsage] _copy_tokenizer, # pyright: ignore[reportPrivateUsage] _find_snapshot, # pyright: ignore[reportPrivateUsage] _reduce_config, # pyright: ignore[reportPrivateUsage] ) def _make_task( content: str = "Hello, what is 2+2?", max_tokens: int = 10, seed: int = 42, ) -> TextGenerationTaskParams: return TextGenerationTaskParams( model=ModelId("test"), input=[InputMessage(role="user", content=content)], max_output_tokens=max_tokens, temperature=0.7, seed=seed, ) # ── Helpers ──────────────────────────────────────────────────────────────── # def _collect_mlx_generate( model: Model, tokenizer: TokenizerWrapper, task: TextGenerationTaskParams, kv_prefix_cache: KVPrefixCache | None, ) -> list[int]: """Run mlx_generate and collect output token IDs.""" prompt = apply_chat_template(tokenizer=tokenizer, task_params=task) tokens: list[int] = [] for resp in mlx_generate( model=model, tokenizer=tokenizer, task=task, prompt=prompt, kv_prefix_cache=kv_prefix_cache, group=None, ): tokens.append(resp.token) if resp.finish_reason is not None: break return tokens def _collect_batch_generate( model: Model, tokenizer: TokenizerWrapper, task_params: TextGenerationTaskParams, kv_prefix_cache: KVPrefixCache | None, ) -> list[int]: """Run ExoBatchGenerator and collect raw output token IDs""" exo_gen = ExoBatchGenerator( model=model, tokenizer=tokenizer, group=None, kv_prefix_cache=kv_prefix_cache, ) prompt = apply_chat_template(tokenizer=tokenizer, task_params=task_params) exo_gen.submit(task_params=task_params, prompt=prompt) tokens: list[int] = [] while exo_gen.has_work: results = exo_gen.step() for _uid, response in results: tokens.append(response.token) exo_gen.close() return tokens def _assert_state_equal(sa: object, sb: object, label: str) -> None: """Compare two state items, handling both plain arrays and tuples of arrays (CacheList).""" if isinstance(sa, tuple): assert isinstance(sb, tuple), f"{label}: type mismatch" for k, (arr_a, arr_b) in enumerate( zip( cast(tuple[mx.array, ...], sa), cast(tuple[mx.array, ...], sb), strict=True, ) ): a_f = mx.array(arr_a).astype(mx.float32) b_f = mx.array(arr_b).astype(mx.float32) if a_f.size == 0: assert b_f.size == 0, f"{label}[{k}]: size mismatch" continue diff = float(mx.max(mx.abs(a_f - b_f)).item()) assert diff == 0.0, f"{label}[{k}]: max diff {diff}" else: sa_f = mx.array(cast(mx.array, sa)).astype(mx.float32) sb_f = mx.array(cast(mx.array, sb)).astype(mx.float32) if sa_f.size == 0: assert sb_f.size == 0, f"{label}: size mismatch" return diff = float(mx.max(mx.abs(sa_f - sb_f)).item()) assert diff == 0.0, f"{label}: max diff {diff}" def _compare_cache_arrays( cache_a: KVCacheType, cache_b: KVCacheType, label: str = "", ) -> None: """Assert two KV caches have identical array values.""" assert len(cache_a) == len(cache_b), ( f"{label}Cache layer count: {len(cache_a)} vs {len(cache_b)}" ) for i, (a, b) in enumerate(zip(cache_a, cache_b, strict=True)): assert type(a) is type(b), ( f"{label}Layer {i}: type {type(a).__name__} vs {type(b).__name__}" ) states_a = a.state states_b = b.state assert len(states_a) == len(states_b), ( f"{label}Layer {i}: state count {len(states_a)} vs {len(states_b)}" ) for j, (sa, sb) in enumerate(zip(states_a, states_b, strict=True)): if sa is None and sb is None: continue assert sa is not None and sb is not None, ( f"{label}Layer {i}, state {j}: one is None" ) _assert_state_equal(sa, sb, f"{label}Layer {i}, state {j}") def _safe_state(cache: object) -> list[object]: """Safely access .state on a cache object. Returns [] if uninitialized.""" # RotatingKVCache.state crashes when keys is None (uninitialized) if getattr(cache, "keys", _SENTINEL) is None: return [] try: return list(cache.state) # type: ignore[union-attr] except (AttributeError, TypeError): return [] _SENTINEL = object() def _compare_snapshots( snaps_a: list[CacheSnapshot] | None, snaps_b: list[CacheSnapshot] | None, label: str = "", ) -> None: """Assert two snapshot lists are identical.""" if snaps_a is None: assert snaps_b is None, f"{label}One side has snapshots, other doesn't" return assert snaps_b is not None, f"{label}One side has snapshots, other doesn't" assert len(snaps_a) == len(snaps_b), ( f"{label}Snapshot count: {len(snaps_a)} vs {len(snaps_b)}" ) for k, (sa, sb) in enumerate(zip(snaps_a, snaps_b, strict=True)): assert sa.token_count == sb.token_count, ( f"{label}Snapshot {k} token_count: {sa.token_count} vs {sb.token_count}" ) for layer_i, (s1, s2) in enumerate(zip(sa.states, sb.states, strict=True)): if s1 is None and s2 is None: continue assert s1 is not None and s2 is not None, ( f"{label}Snapshot {k}, layer {layer_i}: one state is None" ) state_a = _safe_state(s1) state_b = _safe_state(s2) if not state_a and not state_b: continue assert len(state_a) == len(state_b), ( f"{label}Snapshot {k}, layer {layer_i}: state length mismatch" ) for st_j, (arr_a, arr_b) in enumerate(zip(state_a, state_b, strict=True)): if arr_a is None and arr_b is None: continue assert arr_a is not None and arr_b is not None _assert_state_equal( arr_a, arr_b, f"{label}Snapshot {k}, layer {layer_i}, state {st_j}", ) # ── Test class ────────────────────────────────────────────────────────────── # @pytest.mark.slow class TestBatchVsGenerate: """Verify BatchGenerator matches mlx_generate for output tokens and prefix cache.""" @pytest.fixture(autouse=True) def _cleanup(self): yield mx.clear_cache() gc.collect() @pytest.mark.parametrize( "spec", ARCHITECTURES, ids=[a.name for a in ARCHITECTURES], ) def test_same_output_and_cache(self, spec: ArchSpec) -> None: if not _arch_available(spec): pytest.skip(f"Model {spec.hub_name} not cached locally") snapshot = _find_snapshot(spec.hub_name) assert snapshot is not None tmpdir = Path(tempfile.mkdtemp(prefix=f"exo_batchtest_{spec.name}_")) try: # Build reduced config with open(snapshot / "config.json") as f: cfg = cast(dict[str, Any], json.load(f)) reduced = _reduce_config(copy.deepcopy(cfg)) (tmpdir / "config.json").write_text(json.dumps(reduced)) # Copy tokenizer tok_src = snapshot if spec.tokenizer_hub is not None: alt = _find_snapshot(spec.tokenizer_hub) if alt is not None: tok_src = alt _copy_tokenizer(tok_src, tmpdir) # Load tokenizer, build model with random weights model_id = ModelId(f"mlx-community/{spec.hub_name}") tokenizer = load_tokenizer_for_model_id(model_id, tmpdir) mx.random.seed(0) model = _build_model(spec.module, reduced) task = _make_task() # ── Run mlx_generate path ── # Seed is set inside mlx_generate/ExoBatchGenerator.submit from task.seed kv_mlx = KVPrefixCache(None) mlx_tokens = _collect_mlx_generate(model, tokenizer, task, kv_mlx) # ── Run batch generator path ── kv_batch = KVPrefixCache(None) batch_tokens = _collect_batch_generate(model, tokenizer, task, kv_batch) # ── Compare output tokens ── assert len(mlx_tokens) > 0, "mlx_generate produced no tokens" assert len(batch_tokens) > 0, "BatchGenerator produced no tokens" assert mlx_tokens == batch_tokens, ( f"[{spec.name}] Token mismatch:\n" f" mlx_generate: {mlx_tokens}\n" f" BatchGenerator: {batch_tokens}" ) # ── Compare prefix cache KV arrays ── assert len(kv_mlx.caches) == 1, "mlx_generate didn't save to prefix cache" assert len(kv_batch.caches) == 1, ( "BatchGenerator didn't save to prefix cache" ) _compare_cache_arrays( kv_mlx.caches[0], kv_batch.caches[0], label=f"[{spec.name}] ", ) # ── Compare cache lengths ── mlx_len = cache_length(kv_mlx.caches[0]) batch_len = cache_length(kv_batch.caches[0]) assert mlx_len == batch_len, ( f"[{spec.name}] Cache length: mlx={mlx_len} vs batch={batch_len}" ) # ── Compare snapshots ── _compare_snapshots( kv_mlx._snapshots[0], # pyright: ignore[reportPrivateUsage] kv_batch._snapshots[0], # pyright: ignore[reportPrivateUsage] label=f"[{spec.name}] ", ) finally: shutil.rmtree(tmpdir, ignore_errors=True) @pytest.mark.parametrize( "spec", ARCHITECTURES, ids=[a.name for a in ARCHITECTURES], ) def test_concurrent_batch_completes(self, spec: ArchSpec) -> None: """Two requests processed concurrently must both complete without crashing and produce non-empty output. Note: batch decode logits are NOT bit-exact with sequential because Metal's matmul kernel picks different reduction tiling for B=1 vs B=2 when L=1 (decode step). This introduces sub-ULP float16 diffs in gate_proj/down_proj/lm_head which swiglu amplifies by |up_values|. With random weights these accumulate into argmax flips; with trained weights the diffs are absorbed and output matches exactly (verified with real Llama-3.2-1B-Instruct-4bit weights). """ if not _arch_available(spec): pytest.skip(f"Model {spec.hub_name} not cached locally") snapshot = _find_snapshot(spec.hub_name) assert snapshot is not None tmpdir = Path(tempfile.mkdtemp(prefix=f"exo_concurrent_{spec.name}_")) try: with open(snapshot / "config.json") as f: cfg = cast(dict[str, Any], json.load(f)) reduced = _reduce_config(copy.deepcopy(cfg)) (tmpdir / "config.json").write_text(json.dumps(reduced)) tok_src = snapshot if spec.tokenizer_hub is not None: alt = _find_snapshot(spec.tokenizer_hub) if alt is not None: tok_src = alt _copy_tokenizer(tok_src, tmpdir) model_id = ModelId(f"mlx-community/{spec.hub_name}") tokenizer = load_tokenizer_for_model_id(model_id, tmpdir) mx.random.seed(0) model = _build_model(spec.module, reduced) # Two different prompts → different prompt lengths. task_a = _make_task(content="Hello, what is 2+2?", seed=42) task_a = task_a.model_copy(update={"temperature": 0.0}) task_b = _make_task( content="Write a short poem about the ocean and the sky.", seed=99, ) task_b = task_b.model_copy(update={"temperature": 0.0}) # ── Concurrent: submit both to one ExoBatchGenerator ── exo_gen = ExoBatchGenerator( model=model, tokenizer=tokenizer, group=None, kv_prefix_cache=None, ) prompt_a = apply_chat_template(tokenizer=tokenizer, task_params=task_a) prompt_b = apply_chat_template(tokenizer=tokenizer, task_params=task_b) uid_a = exo_gen.submit(task_params=task_a, prompt=prompt_a) uid_b = exo_gen.submit(task_params=task_b, prompt=prompt_b) batch_tokens: dict[int, list[int]] = {uid_a: [], uid_b: []} finished: set[int] = set() while exo_gen.has_work: results = exo_gen.step() for uid, response in results: batch_tokens[uid].append(response.token) if response.finish_reason is not None: finished.add(uid) exo_gen.close() # ── Verify both completed ── assert len(batch_tokens[uid_a]) > 0, "No tokens for task A" assert len(batch_tokens[uid_b]) > 0, "No tokens for task B" assert uid_a in finished, "Task A never finished" assert uid_b in finished, "Task B never finished" finally: shutil.rmtree(tmpdir, ignore_errors=True) ================================================ FILE: src/exo/worker/tests/unittests/test_mlx/test_distributed_fix.py ================================================ import multiprocessing as mp import os from dataclasses import dataclass from typing import Any, Callable import pytest from exo.worker.tests.unittests.test_mlx.conftest import ( DEFAULT_GPT_OSS_CONFIG, create_hostfile, run_gpt_oss_pipeline_device, run_gpt_oss_tensor_parallel_device, ) def _check_model_exists() -> bool: return DEFAULT_GPT_OSS_CONFIG.model_path.exists() pytestmark = [ pytest.mark.slow, pytest.mark.skipif( not _check_model_exists(), reason=f"GPT-OSS model not found at {DEFAULT_GPT_OSS_CONFIG.model_path}", ), ] @dataclass class DistributedTestResult: timed_out: bool world_size: int results: dict[int, tuple[bool, str]] @property def all_success(self) -> bool: if len(self.results) != self.world_size: return False return all(r[0] for r in self.results.values()) def run_distributed_test( world_size: int, port_offset: int, process_timeout: int, target: Callable[..., None], make_args: Callable[[int], tuple[Any, ...]], ) -> DistributedTestResult: ctx = mp.get_context("spawn") hostfile_path, _ = create_hostfile( world_size, DEFAULT_GPT_OSS_CONFIG.base_port + port_offset ) try: result_queue: Any = ctx.Queue() processes: list[Any] = [] for rank in range(world_size): args = make_args(rank) p = ctx.Process( target=target, args=(rank, world_size, hostfile_path, *args, result_queue), ) p.start() processes.append(p) for p in processes: # pyright: ignore[reportAny] p.join(timeout=process_timeout) # pyright: ignore[reportAny] timed_out = any(p.is_alive() for p in processes) # pyright: ignore[reportAny] for p in processes: # pyright: ignore[reportAny] if p.is_alive(): # pyright: ignore[reportAny] p.terminate() # pyright: ignore[reportAny] p.join(timeout=5) # pyright: ignore[reportAny] results: dict[int, tuple[bool, str]] = {} while not result_queue.empty(): # pyright: ignore[reportAny] rank, success, value = result_queue.get() # pyright: ignore[reportAny] results[rank] = (success, value) return DistributedTestResult( timed_out=timed_out, world_size=world_size, results=results ) finally: os.unlink(hostfile_path) def run_pipeline_test( layer_splits: list[tuple[int, int]], prompt_tokens: int, prefill_step_size: int, port_offset: int = 0, process_timeout: int = 60, ) -> DistributedTestResult: def make_args(rank: int) -> tuple[Any, ...]: return ( layer_splits, prompt_tokens, prefill_step_size, ) return run_distributed_test( world_size=len(layer_splits), port_offset=port_offset, process_timeout=process_timeout, target=run_gpt_oss_pipeline_device, make_args=make_args, ) def run_tensor_test( prompt_tokens: int, prefill_step_size: int, port_offset: int = 0, process_timeout: int = 60, ) -> DistributedTestResult: def make_args(rank: int) -> tuple[Any, ...]: return ( prompt_tokens, prefill_step_size, ) return run_distributed_test( world_size=2, port_offset=port_offset, process_timeout=process_timeout, target=run_gpt_oss_tensor_parallel_device, make_args=make_args, ) class TestPipelineParallelFix: BUG_TRIGGER_SPLITS: list[tuple[int, int]] = [(0, 1), (1, 24)] def test_pipeline_single_layer_first_device(self) -> None: result = run_pipeline_test( layer_splits=self.BUG_TRIGGER_SPLITS, prompt_tokens=100, prefill_step_size=64, process_timeout=60, ) assert not result.timed_out, "Unexpected timeout - fix may not be working" assert result.all_success, f"Failures: {result.results}" class TestPipelineSplitConfigurations: @pytest.mark.parametrize( "layer_splits", [ [(0, 1), (1, 24)], [(0, 6), (6, 24)], [(0, 12), (12, 24)], ], ids=["1_23", "6_18", "12_12"], ) def test_pipeline_splits( self, layer_splits: list[tuple[int, int]], ) -> None: result = run_pipeline_test( layer_splits=layer_splits, prompt_tokens=600, prefill_step_size=512, port_offset=100, ) assert not result.timed_out, f"Timeout with {layer_splits}" assert result.all_success, f"Failures with {layer_splits}: {result.results}" class TestPrefillStepSizeBoundaries: @pytest.mark.parametrize( "prefill_step_size,prompt_tokens", [ (512, 511), (512, 512), (512, 513), (512, 1024), ], ids=["under", "exact", "over", "double"], ) def test_boundary_conditions( self, prefill_step_size: int, prompt_tokens: int, ) -> None: result = run_pipeline_test( layer_splits=[(0, 12), (12, 24)], prompt_tokens=prompt_tokens, prefill_step_size=prefill_step_size, port_offset=200, ) assert not result.timed_out, f"Timeout: {prompt_tokens=}, {prefill_step_size=}" assert result.all_success, f"Failures: {result.results}" class TestTensorParallelFix: def test_tensor_parallel(self) -> None: result = run_tensor_test( prompt_tokens=100, prefill_step_size=64, port_offset=400, ) assert not result.timed_out, "Unexpected timeout" assert result.all_success, f"Failures: {result.results}" class TestTensorParallelBoundaries: @pytest.mark.parametrize( "prefill_step_size,prompt_tokens", [ (512, 511), (512, 512), (512, 513), (512, 1024), ], ids=["under", "exact", "over", "double"], ) def test_tensor_parallel_boundaries( self, prefill_step_size: int, prompt_tokens: int, ) -> None: result = run_tensor_test( prompt_tokens=prompt_tokens, prefill_step_size=prefill_step_size, port_offset=500, ) assert not result.timed_out, f"Timeout: {prompt_tokens=}, {prefill_step_size=}" assert result.all_success, f"Failures: {result.results}" ================================================ FILE: src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py ================================================ # type: ignore import time from typing import cast from unittest.mock import patch import mlx.core as mx import pytest from mlx_lm.models.cache import KVCache from mlx_lm.sample_utils import make_sampler from exo.shared.types.common import ModelId from exo.shared.types.mlx import Model from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams from exo.worker.engines.mlx.cache import ( KVPrefixCache, cache_length, encode_prompt, get_prefix_length, make_kv_cache, ) from exo.worker.engines.mlx.generator.generate import mlx_generate, prefill from exo.worker.engines.mlx.utils_mlx import apply_chat_template from exo.worker.tests.unittests.test_mlx.conftest import ( DEFAULT_GPT_OSS_CONFIG, DEFAULT_GPT_OSS_MODEL_ID, ) def _check_model_exists() -> bool: return DEFAULT_GPT_OSS_CONFIG.model_path.exists() class TestGetPrefixLength: def test_identical_arrays(self): a = mx.array([1, 2, 3, 4, 5]) b = mx.array([1, 2, 3, 4, 5]) assert get_prefix_length(a, b) == 5 def test_no_common_prefix(self): a = mx.array([1, 2, 3]) b = mx.array([4, 5, 6]) assert get_prefix_length(a, b) == 0 def test_partial_prefix(self): a = mx.array([1, 2, 3, 4, 5]) b = mx.array([1, 2, 3, 7, 8]) assert get_prefix_length(a, b) == 3 def test_prompt_longer_than_cached(self): a = mx.array([1, 2, 3, 4, 5]) b = mx.array([1, 2, 3]) assert get_prefix_length(a, b) == 3 def test_cached_longer_than_prompt(self): a = mx.array([1, 2, 3]) b = mx.array([1, 2, 3, 4, 5]) assert get_prefix_length(a, b) == 3 def test_single_token_match(self): a = mx.array([1, 2, 3]) b = mx.array([1, 5, 6]) assert get_prefix_length(a, b) == 1 def test_empty_prompt(self): a = mx.array([]).astype(mx.int32) b = mx.array([1, 2, 3]) assert get_prefix_length(a, b) == 0 def test_empty_cached(self): a = mx.array([1, 2, 3]) b = mx.array([]).astype(mx.int32) assert get_prefix_length(a, b) == 0 def test_both_empty(self): a = mx.array([]).astype(mx.int32) b = mx.array([]).astype(mx.int32) assert get_prefix_length(a, b) == 0 class TestKVPrefix: @pytest.fixture def mock_tokenizer(self): """Create a minimal mock tokenizer for tests that don't need real tokenization.""" from unittest.mock import MagicMock tokenizer = MagicMock() tokenizer.encode.return_value = [1, 2, 3] return tokenizer def test_starts_empty(self, mock_tokenizer): cache = KVPrefixCache(None) assert len(cache.prompts) == 0 assert len(cache.caches) == 0 def test_clear_empties_cache(self, mock_tokenizer): cache = KVPrefixCache(None) cache.prompts.append(mx.array([1, 2, 3])) cache.caches.append([KVCache()]) cache.clear() assert len(cache.prompts) == 0 assert len(cache.caches) == 0 def test_clear_on_empty_cache(self, mock_tokenizer): cache = KVPrefixCache(None) cache.clear() assert len(cache.prompts) == 0 def _load_gpt_oss() -> tuple[Model, object]: from mlx_lm.utils import load_model from exo.worker.engines.mlx.utils_mlx import load_tokenizer_for_model_id model_path = DEFAULT_GPT_OSS_CONFIG.model_path model_id = ModelId(DEFAULT_GPT_OSS_MODEL_ID) model, _ = load_model(model_path, lazy=False) tokenizer = load_tokenizer_for_model_id(model_id, model_path) return cast(Model, model), tokenizer @pytest.mark.slow @pytest.mark.skipif( not _check_model_exists(), reason=f"GPT-OSS model not found at {DEFAULT_GPT_OSS_CONFIG.model_path}", ) class TestKVPrefixCacheWithModel: @pytest.fixture(scope="class") def model_and_tokenizer(self): model, tokenizer = _load_gpt_oss() return model, tokenizer def test_prefill_populates_cache(self, model_and_tokenizer): model, tokenizer = model_and_tokenizer task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content="Hello!!")], max_output_tokens=1, ) prompt = apply_chat_template(tokenizer, task) tokens = encode_prompt(tokenizer, prompt) cache = make_kv_cache(model) _, _, snapshots = prefill( model, tokenizer, make_sampler(0.0), tokens, cache, group=None, on_prefill_progress=None, distributed_prompt_progress_callback=None, ) # Cache should now hold the prompt tokens minus one assert cache_length(cache) == len(tokens) - 1 # Snapshots should be available for models with non-KV caches assert len(snapshots) > 0 def test_add_and_get_exact_match(self, model_and_tokenizer): model, tokenizer = model_and_tokenizer task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content="Test exact")], max_output_tokens=1, ) prompt = apply_chat_template(tokenizer, task) tokens = encode_prompt(tokenizer, prompt) cache = make_kv_cache(model) _, _, snapshots = prefill( model, tokenizer, make_sampler(0.0), tokens, cache, group=None, on_prefill_progress=None, distributed_prompt_progress_callback=None, ) kv_prefix_cache = KVPrefixCache(None) kv_prefix_cache.add_kv_cache(tokens, cache, snapshots) assert len(kv_prefix_cache.prompts) == 1 stored_length = cache_length(kv_prefix_cache.caches[0]) assert stored_length > 0 # Retrieve with same prompt: exact match result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache( model, tokens ) assert matched_index == 0 # Exact match returns last token(s) — for models with SSM/rotating caches, # snapshot availability constrains how far back we can trim, so remaining # may be 1 or 2 tokens depending on the model. assert len(remaining_tokens) >= 1 assert mx.array_equal(remaining_tokens, tokens[-len(remaining_tokens) :]) def test_add_and_get_prefix_match(self, model_and_tokenizer): """get_kv_cache with a longer prompt sharing prefix should return partial match.""" model, tokenizer = model_and_tokenizer short_task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content="Hi")], max_output_tokens=1, ) short_prompt = apply_chat_template(tokenizer, short_task) short_tokens = encode_prompt(tokenizer, short_prompt) cache = make_kv_cache(model) _, _, snapshots = prefill( model, tokenizer, make_sampler(0.0), short_tokens, cache, group=None, on_prefill_progress=None, distributed_prompt_progress_callback=None, ) kv_prefix_cache = KVPrefixCache(None) kv_prefix_cache.add_kv_cache(short_tokens, cache, snapshots) # Query with longer prompt that shares the chat template prefix long_task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content="Hi there, how are you?")], max_output_tokens=1, ) long_prompt = apply_chat_template(tokenizer, long_task) long_tokens = encode_prompt(tokenizer, long_prompt) # The prompts share a prefix (chat template preamble + "Hi") expected_prefix = get_prefix_length(long_tokens, short_tokens) assert expected_prefix > 0, ( "Prompts should share a prefix from the chat template" ) result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache( model, long_tokens ) assert matched_index == 0 # remaining_tokens covers from snapshot restore position to end assert len(remaining_tokens) >= len(long_tokens) - expected_prefix def test_stored_cache_not_mutated_after_get_and_generation( self, model_and_tokenizer ): """Getting a cache and then mutating it (as generation does) must not corrupt stored cache.""" model, tokenizer = model_and_tokenizer task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content="Mutation test")], max_output_tokens=1, ) prompt = apply_chat_template(tokenizer, task) tokens = encode_prompt(tokenizer, prompt) cache = make_kv_cache(model) _, _, snapshots = prefill( model, tokenizer, make_sampler(0.0), tokens, cache, group=None, on_prefill_progress=None, distributed_prompt_progress_callback=None, ) kv_prefix_cache = KVPrefixCache(None) kv_prefix_cache.add_kv_cache(tokens, cache, snapshots) stored_length = cache_length(kv_prefix_cache.caches[0]) # Get cache and mutate it (simulating what generation does) result_cache, _, matched_index = kv_prefix_cache.get_kv_cache(model, tokens) assert matched_index == 0 # Simulate generation: feed many additional tokens through the cache head_dim = result_cache[0].keys.shape[-1] num_heads = result_cache[0].keys.shape[1] extra_keys = mx.random.normal((1, num_heads, 50, head_dim)) extra_values = mx.random.normal((1, num_heads, 50, head_dim)) for layer_cache in result_cache: layer_cache.update_and_fetch(extra_keys, extra_values) mx.eval([c.keys for c in result_cache]) # Stored cache must be unchanged assert cache_length(kv_prefix_cache.caches[0]) == stored_length def test_stored_cache_survives_repeated_get_mutate_cycles( self, model_and_tokenizer ): """Multiple get+mutate cycles (like repeated user requests) must not corrupt cache.""" model, tokenizer = model_and_tokenizer task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content="Repeat test")], max_output_tokens=1, ) prompt = apply_chat_template(tokenizer, task) tokens = encode_prompt(tokenizer, prompt) cache = make_kv_cache(model) _, _, snapshots = prefill( model, tokenizer, make_sampler(0.0), tokens, cache, group=None, on_prefill_progress=None, distributed_prompt_progress_callback=None, ) kv_prefix_cache = KVPrefixCache(None) kv_prefix_cache.add_kv_cache(tokens, cache, snapshots) stored_length = cache_length(kv_prefix_cache.caches[0]) for i in range(3): result_cache, _, _ = kv_prefix_cache.get_kv_cache(model, tokens) head_dim = result_cache[0].keys.shape[-1] num_heads = result_cache[0].keys.shape[1] extra = mx.random.normal((1, num_heads, 30, head_dim)) for layer_cache in result_cache: layer_cache.update_and_fetch(extra, extra) mx.eval([c.keys for c in result_cache]) assert cache_length(kv_prefix_cache.caches[0]) == stored_length, ( f"Failed on loop {i}" ) def test_mlx_generate_populates_cache(self, model_and_tokenizer): """mlx_generate should save the cache after generation completes.""" model, tokenizer = model_and_tokenizer kv_prefix_cache = KVPrefixCache(None) task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content="Hello")], max_output_tokens=5, ) prompt = apply_chat_template(tokenizer, task) prompt_tokens = encode_prompt(tokenizer, prompt) # Consume the entire generator so the cache-saving code after yield runs generated_tokens = 0 for _response in mlx_generate( model=model, tokenizer=tokenizer, task=task, prompt=prompt, kv_prefix_cache=kv_prefix_cache, group=None, ): generated_tokens += 1 assert len(kv_prefix_cache.prompts) == 1 assert len(kv_prefix_cache.caches) == 1 # Cache should contain prompt + generated tokens expected_length = len(prompt_tokens) + generated_tokens assert cache_length(kv_prefix_cache.caches[0]) == expected_length def test_mlx_generate_second_call_gets_prefix_hit(self, model_and_tokenizer): """Second mlx_generate call with same prompt should get a prefix hit from stored cache.""" model, tokenizer = model_and_tokenizer kv_prefix_cache = KVPrefixCache(None) task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content="Reuse test")], max_output_tokens=5, ) prompt = apply_chat_template(tokenizer, task) prompt_tokens = encode_prompt(tokenizer, prompt) # First generation populates cache for _response in mlx_generate( model=model, tokenizer=tokenizer, task=task, prompt=prompt, kv_prefix_cache=kv_prefix_cache, group=None, ): pass assert len(kv_prefix_cache.prompts) == 1 # Second call should find a prefix match (the stored cache contains # prompt + generated tokens, which shares the prompt prefix) result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache( model, prompt_tokens ) # The stored cache is longer than the prompt (it includes generated tokens), # so this is a prefix match where our prompt is fully contained assert matched_index == 0 # Exact match: remaining_tokens is just the last token and the one before assert len(remaining_tokens) == 2 assert mx.array_equal(remaining_tokens, prompt_tokens[-2:]) def test_mlx_generate_long_prompt_updates_cache_in_place(self, model_and_tokenizer): """With a prompt > 1000 tokens, second generation should update the cache entry in-place.""" model, tokenizer = model_and_tokenizer kv_prefix_cache = KVPrefixCache(None) # Build a long user message (> 1000 tokens) to exceed _MIN_PREFIX_HIT_TO_UPDATE base_text = "The quick brown fox jumps over the lazy dog. " base_tokens = tokenizer.encode(base_text) repeats = (1200 // len(base_tokens)) + 2 long_content = base_text * repeats task1 = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content=long_content)], max_output_tokens=5, ) prompt1 = apply_chat_template(tokenizer, task1) prompt1_tokens = encode_prompt(tokenizer, prompt1) assert len(prompt1_tokens) > 1000, ( "Prompt must exceed _MIN_PREFIX_HIT_TO_UPDATE" ) # First generation populates the cache (must prefill all tokens) t0 = time.perf_counter() for _response in mlx_generate( model=model, tokenizer=tokenizer, task=task1, prompt=prompt1, kv_prefix_cache=kv_prefix_cache, group=None, ): pass first_gen_time = time.perf_counter() - t0 assert len(kv_prefix_cache.prompts) == 1 first_cache_length = cache_length(kv_prefix_cache.caches[0]) # Second generation: same long prompt + extra content (simulating multi-turn) task2 = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[ InputMessage(role="user", content=long_content), InputMessage(role="assistant", content="Sure, I can help."), InputMessage(role="user", content="Tell me more."), ], max_output_tokens=5, ) prompt2 = apply_chat_template(tokenizer, task2) prompt2_tokens = encode_prompt(tokenizer, prompt2) # Verify the prompts share a long prefix prefix_len = get_prefix_length(prompt2_tokens, prompt1_tokens) assert prefix_len > 1000, "Prompts must share > 1000 token prefix" # Second generation should reuse the cached prefix (only prefill new tokens) t0 = time.perf_counter() for _response in mlx_generate( model=model, tokenizer=tokenizer, task=task2, prompt=prompt2, kv_prefix_cache=kv_prefix_cache, group=None, ): pass second_gen_time = time.perf_counter() - t0 # Second generation should be significantly faster due to prefix cache hit - hopefully not flaky assert second_gen_time < first_gen_time * 0.5, ( f"Expected prefix cache speedup: " f"first={first_gen_time:.2f}s, second={second_gen_time:.2f}s" ) # With prefix_hit > 1000, should update in-place (not add a second entry) assert len(kv_prefix_cache.prompts) == 1 # Updated cache should be longer (prompt2 + generated > prompt1 + generated) updated_cache_length = cache_length(kv_prefix_cache.caches[0]) assert updated_cache_length > first_cache_length def test_mlx_generate_stored_cache_not_mutated(self, model_and_tokenizer): """After mlx_generate saves a cache, a second generation must not corrupt the stored copy.""" model, tokenizer = model_and_tokenizer kv_prefix_cache = KVPrefixCache(None) task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content="Immutable test")], max_output_tokens=5, ) prompt = apply_chat_template(tokenizer, task) # First generation populates cache for _response in mlx_generate( model=model, tokenizer=tokenizer, task=task, prompt=prompt, kv_prefix_cache=kv_prefix_cache, group=None, ): pass firstcache_length = cache_length(kv_prefix_cache.caches[0]) # Second generation gets the cache and mutates it during generation for _response in mlx_generate( model=model, tokenizer=tokenizer, task=task, prompt=prompt, kv_prefix_cache=kv_prefix_cache, group=None, ): pass # The first stored cache must not have been mutated by the second generation assert cache_length(kv_prefix_cache.caches[0]) == firstcache_length def test_evicts_lru_entry_under_memory_pressure(self, model_and_tokenizer): """Under memory pressure, adding a new cache entry evicts the least recently used one.""" model, tokenizer = model_and_tokenizer kv_prefix_cache = KVPrefixCache(None) # Add three cache entries with different prompts prompts = ["First entry", "Second entry", "Third entry"] for i, content in enumerate(prompts): task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content=content)], max_output_tokens=1, ) prompt = apply_chat_template(tokenizer, task) tokens = encode_prompt(tokenizer, prompt) cache = make_kv_cache(model) prefill( model, tokenizer, make_sampler(0.0), tokens, cache, group=None, on_prefill_progress=None, distributed_prompt_progress_callback=None, ) kv_prefix_cache.add_kv_cache(tokens, cache) # Stagger _last_used so LRU order is deterministic kv_prefix_cache._last_used[i] = float(i) assert len(kv_prefix_cache.prompts) == 3 # Access the third entry to make it most recently used kv_prefix_cache._last_used[2] = 100.0 # Entry 0 (_last_used=0.0) is LRU, entry 1 (_last_used=1.0) is next # Simulate memory pressure: return usage above _MEMORY_THRESHOLD (0.9) with patch( "exo.worker.engines.mlx.cache.get_memory_used_percentage", return_value=0.95, ): # Trigger eviction by adding a new entry task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content="New entry")], max_output_tokens=1, ) prompt = apply_chat_template(tokenizer, task) tokens = encode_prompt(tokenizer, prompt) cache = make_kv_cache(model) prefill( model, tokenizer, make_sampler(0.0), tokens, cache, group=None, on_prefill_progress=None, distributed_prompt_progress_callback=None, ) kv_prefix_cache.add_kv_cache(tokens, cache) # LRU entries should have been evicted (entries 0, 1, 2 in order of _last_used) # Since fake_active stays above threshold after each eviction (we don't change it), # all old entries get evicted, leaving only the newly added one assert len(kv_prefix_cache.prompts) == 1 # The surviving entry should be the newly added one assert get_prefix_length(kv_prefix_cache.prompts[0], tokens) == len(tokens) ================================================ FILE: src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py ================================================ # type: ignore """Test that pipeline prefill callbacks and output exactly match stream_generate. Spins up a single-device (non-pipeline) run and a distributed pipeline run, then verifies that the prompt_progress_callback sequences are identical and that generated text matches. """ import json import multiprocessing as mp import os import tempfile import traceback from typing import Any, cast import pytest from exo.shared.constants import EXO_MODELS_DIR from exo.shared.models.model_cards import ModelCard, ModelTask from exo.shared.types.common import ModelId from exo.shared.types.memory import Memory from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams MODEL_ID = "mlx-community/gpt-oss-20b-MXFP4-Q8" MODEL_PATH = EXO_MODELS_DIR / "mlx-community--gpt-oss-20b-MXFP4-Q8" TOTAL_LAYERS = 24 MAX_TOKENS = 10 SEED = 42 TEMPERATURE = 0.0 def _model_card() -> ModelCard: return ModelCard( model_id=ModelId(MODEL_ID), storage_size=Memory.from_gb(12), n_layers=TOTAL_LAYERS, hidden_size=2880, supports_tensor=False, tasks=[ModelTask.TextGeneration], ) def _build_prompt(tokenizer: Any, prompt_tokens: int) -> tuple[str, Any]: """Build a prompt with the given number of user-content tokens, return (chat_prompt, task).""" from exo.worker.engines.mlx.utils_mlx import apply_chat_template base_text = "The quick brown fox jumps over the lazy dog. " base_toks = tokenizer.encode(base_text) repeats = (prompt_tokens // len(base_toks)) + 2 long_text = base_text * repeats tokens = tokenizer.encode(long_text)[:prompt_tokens] prompt_text = tokenizer.decode(tokens) task = TextGenerationTaskParams( model=MODEL_ID, input=[InputMessage(role="user", content=prompt_text)], max_output_tokens=MAX_TOKENS, temperature=TEMPERATURE, seed=SEED, ) prompt = apply_chat_template(tokenizer, task) return prompt, task # --------------------------------------------------------------------------- # Single-device process: uses stream_generate path (no pipeline layers) # --------------------------------------------------------------------------- def _run_single_device( prompt_tokens: int, result_queue: Any, ) -> None: """Load full model without pipeline sharding, run mlx_generate, record callbacks.""" try: import mlx.core as mx from mlx_lm.utils import load_model from exo.shared.types.worker.shards import PipelineShardMetadata from exo.worker.engines.mlx.cache import encode_prompt from exo.worker.engines.mlx.generator.generate import mlx_generate from exo.worker.engines.mlx.utils_mlx import ( build_model_path, get_tokenizer, ) model_path = build_model_path(ModelId(MODEL_ID)) model, _ = load_model(model_path, lazy=True, strict=False) mx.eval(model) # Use PipelineShardMetadata just for get_tokenizer (needs model_card), but # do NOT apply pipeline sharding — the model keeps all layers unwrapped. dummy_meta = PipelineShardMetadata( model_card=_model_card(), device_rank=0, world_size=1, start_layer=0, end_layer=TOTAL_LAYERS, n_layers=TOTAL_LAYERS, ) tokenizer = get_tokenizer(model_path, dummy_meta) prompt, task = _build_prompt(tokenizer, prompt_tokens) callbacks: list[tuple[int, int]] = [] def on_progress(processed: int, total: int) -> None: callbacks.append((processed, total)) generated_text = "" for response in mlx_generate( model=model, tokenizer=tokenizer, task=task, prompt=prompt, kv_prefix_cache=None, group=None, on_prefill_progress=on_progress, ): generated_text += response.text if response.finish_reason is not None: break # Also record the token count that prefill() received (prompt_tokens[:-1]) all_tokens = encode_prompt(tokenizer, prompt) prefill_token_count = len(all_tokens) - 1 result_queue.put( ( True, { "callbacks": callbacks, "text": generated_text, "prefill_token_count": prefill_token_count, }, ) ) except Exception as e: result_queue.put((False, f"{e}\n{traceback.format_exc()}")) # --------------------------------------------------------------------------- # Pipeline device process: uses _pipeline_prefill_cache path # --------------------------------------------------------------------------- def _run_pipeline_device( rank: int, world_size: int, hostfile_path: str, layer_splits: list[tuple[int, int]], prompt_tokens: int, result_queue: Any, ) -> None: """Load model with pipeline sharding, run mlx_generate, record callbacks.""" os.environ["MLX_HOSTFILE"] = hostfile_path os.environ["MLX_RANK"] = str(rank) try: import mlx.core as mx from exo.shared.types.worker.shards import PipelineShardMetadata from exo.worker.engines.mlx.cache import encode_prompt from exo.worker.engines.mlx.generator.generate import mlx_generate from exo.worker.engines.mlx.utils_mlx import shard_and_load group = mx.distributed.init(backend="ring", strict=True) start_layer, end_layer = layer_splits[rank] shard_meta = PipelineShardMetadata( model_card=_model_card(), device_rank=rank, world_size=world_size, start_layer=start_layer, end_layer=end_layer, n_layers=TOTAL_LAYERS, ) model, tokenizer = shard_and_load( shard_meta, group, on_timeout=None, on_layer_loaded=None ) model = cast(Any, model) prompt, task = _build_prompt(tokenizer, prompt_tokens) callbacks: list[tuple[int, int]] = [] def on_progress(processed: int, total: int) -> None: callbacks.append((processed, total)) def distributed_prompt_progress_callback(_group: Any = group) -> None: from exo.worker.engines.mlx.utils_mlx import mx_any mx_any(False, _group) generated_text = "" for response in mlx_generate( model=model, tokenizer=tokenizer, task=task, prompt=prompt, kv_prefix_cache=None, group=group, on_prefill_progress=on_progress, distributed_prompt_progress_callback=distributed_prompt_progress_callback, ): generated_text += response.text if response.finish_reason is not None: break all_tokens = encode_prompt(tokenizer, prompt) prefill_token_count = len(all_tokens) - 1 result_queue.put( ( rank, True, { "callbacks": callbacks, "text": generated_text, "prefill_token_count": prefill_token_count, }, ) ) except Exception as e: result_queue.put((rank, False, f"{e}\n{traceback.format_exc()}")) # --------------------------------------------------------------------------- # Test helpers # --------------------------------------------------------------------------- def _create_hostfile(world_size: int, base_port: int) -> str: hosts = [f"127.0.0.1:{base_port + i}" for i in range(world_size)] with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: json.dump(hosts, f) return f.name def _run_single_device_test(prompt_tokens: int, timeout: int = 120) -> dict[str, Any]: """Run single-device (stream_generate) prefill and return results.""" ctx = mp.get_context("spawn") result_queue: Any = ctx.Queue() p = ctx.Process(target=_run_single_device, args=(prompt_tokens, result_queue)) p.start() p.join(timeout=timeout) if p.is_alive(): p.terminate() p.join(timeout=5) pytest.fail("Single-device process timed out") assert not result_queue.empty(), "Single-device process produced no result" success, data = result_queue.get() assert success, f"Single-device process failed:\n{data}" return data def _run_pipeline_test( layer_splits: list[tuple[int, int]], prompt_tokens: int, base_port: int, timeout: int = 120, ) -> dict[int, dict[str, Any]]: """Run pipeline prefill across ranks and return per-rank results.""" world_size = len(layer_splits) hostfile_path = _create_hostfile(world_size, base_port) ctx = mp.get_context("spawn") result_queue: Any = ctx.Queue() try: processes: list[Any] = [] for rank in range(world_size): p = ctx.Process( target=_run_pipeline_device, args=( rank, world_size, hostfile_path, layer_splits, prompt_tokens, result_queue, ), ) p.start() processes.append(p) for p in processes: p.join(timeout=timeout) timed_out = any(p.is_alive() for p in processes) for p in processes: if p.is_alive(): p.terminate() p.join(timeout=5) assert not timed_out, "Pipeline processes timed out" results: dict[int, dict[str, Any]] = {} while not result_queue.empty(): rank, success, data = result_queue.get() assert success, f"Pipeline rank {rank} failed:\n{data}" results[rank] = data assert len(results) == world_size, ( f"Expected {world_size} results, got {len(results)}: missing ranks {set(range(world_size)) - results.keys()}" ) return results finally: os.unlink(hostfile_path) # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- pytestmark = [ pytest.mark.slow, pytest.mark.skipif( not MODEL_PATH.exists(), reason=f"GPT-OSS model not found at {MODEL_PATH}", ), ] LAYER_SPLITS_4WAY: list[tuple[int, int]] = [(0, 6), (6, 12), (12, 18), (18, 24)] LAYER_SPLITS_2WAY: list[tuple[int, int]] = [(0, 12), (12, 24)] class TestPipelineNoDeadlock: """Pipeline prefill must not deadlock at any rank count or prompt length.""" @pytest.mark.parametrize( "layer_splits,prompt_tokens", [ (LAYER_SPLITS_2WAY, 128), (LAYER_SPLITS_2WAY, 4096), (LAYER_SPLITS_2WAY, 8192), (LAYER_SPLITS_2WAY, 16384), (LAYER_SPLITS_4WAY, 128), (LAYER_SPLITS_4WAY, 4096), (LAYER_SPLITS_4WAY, 8192), (LAYER_SPLITS_4WAY, 16384), ], ids=[ "2rank_128tok", "2rank_4096tok", "2rank_8192tok", "2rank_16384tok", "4rank_128tok", "4rank_4096tok", "4rank_8192tok", "4rank_16384tok", ], ) def test_no_deadlock( self, layer_splits: list[tuple[int, int]], prompt_tokens: int, ) -> None: """Pipeline must complete without deadlock at various prompt lengths.""" pipeline_results = _run_pipeline_test( layer_splits=layer_splits, prompt_tokens=prompt_tokens, base_port=29650, timeout=60, ) # If we get here, no deadlock. Verify all ranks produced output. for rank, pipe_data in sorted(pipeline_results.items()): assert pipe_data["text"], f"Rank {rank} produced no output text" class TestPipelinePrefillCallbacks: """Verify that pipeline prefill callbacks exactly match stream_generate callbacks.""" @pytest.mark.parametrize( "prompt_tokens", [50, 500, 5000], ids=["short_50", "medium_500", "long_5000"], ) def test_callbacks_match(self, prompt_tokens: int) -> None: """All pipeline ranks must produce identical callback sequences.""" # Run 4-rank pipeline pipeline_results = _run_pipeline_test( layer_splits=LAYER_SPLITS_4WAY, prompt_tokens=prompt_tokens, base_port=29700, timeout=180, ) # All ranks must agree on prefill token count and callback sequence rank0_data = pipeline_results[0] rank0_callbacks = rank0_data["callbacks"] prefill_count = rank0_data["prefill_token_count"] for rank, pipe_data in sorted(pipeline_results.items()): pipe_callbacks = pipe_data["callbacks"] assert pipe_data["prefill_token_count"] == prefill_count, ( f"Rank {rank} prefill token count mismatch: " f"{pipe_data['prefill_token_count']} vs {prefill_count}" ) assert pipe_callbacks == rank0_callbacks, ( f"Rank {rank} callback mismatch for {prompt_tokens} prompt tokens " f"(prefill M={prefill_count}):\n" f" pipeline R0 ({len(rank0_callbacks)} callbacks): {rank0_callbacks}\n" f" pipeline R{rank} ({len(pipe_callbacks)} callbacks): {pipe_callbacks}" ) # Structural checks: starts with (0, M), ends with (M, M), monotonically increasing assert rank0_callbacks[0] == (0, prefill_count), ( f"First callback should be (0, {prefill_count}), got {rank0_callbacks[0]}" ) assert rank0_callbacks[-1] == (prefill_count, prefill_count), ( f"Last callback should be ({prefill_count}, {prefill_count}), got {rank0_callbacks[-1]}" ) for i in range(1, len(rank0_callbacks)): assert rank0_callbacks[i][0] >= rank0_callbacks[i - 1][0], ( f"Callbacks not monotonically increasing at index {i}: {rank0_callbacks}" ) @pytest.mark.parametrize( "prompt_tokens", [50, 500], ids=["short_50", "medium_500"], ) def test_output_matches(self, prompt_tokens: int) -> None: """Pipeline-generated text must match single-device output.""" single = _run_single_device_test(prompt_tokens, timeout=180) pipeline_results = _run_pipeline_test( layer_splits=LAYER_SPLITS_4WAY, prompt_tokens=prompt_tokens, base_port=29800, timeout=180, ) single_text = single["text"] # The last rank produces the final logits, so its output should match. # Due to SDPA tiling non-determinism, allow minor differences in text. last_rank = max(pipeline_results.keys()) pipe_text = pipeline_results[last_rank]["text"] # For deterministic sampling (temp=0.0), outputs should match exactly # or be very close. Log both for debugging even if they match. if single_text != pipe_text: # Find first divergence point min_len = min(len(single_text), len(pipe_text)) diverge_idx = next( (i for i in range(min_len) if single_text[i] != pipe_text[i]), min_len, ) pytest.fail( f"Output text diverged at character {diverge_idx} for {prompt_tokens} prompt tokens:\n" f" single-device: {single_text!r}\n" f" pipeline R{last_rank}: {pipe_text!r}" ) class TestPipelineCallbacksStructure: """Verify structural properties of callbacks independent of model output.""" def test_callback_structure_matches_generate_step(self) -> None: """Verify callbacks follow generate_step's pattern: (0,M), chunks up to M-1, (M,M).""" prompt_tokens = 200 pipeline_results = _run_pipeline_test( layer_splits=LAYER_SPLITS_4WAY, prompt_tokens=prompt_tokens, base_port=29900, timeout=180, ) for rank, pipe_data in sorted(pipeline_results.items()): callbacks = pipe_data["callbacks"] m = pipe_data["prefill_token_count"] assert m > 0, f"Rank {rank}: prefill token count is 0" assert callbacks[0] == (0, m), ( f"Rank {rank}: first callback should be (0, {m}), got {callbacks[0]}" ) assert callbacks[-1] == (m, m), ( f"Rank {rank}: last callback should be ({m}, {m}), got {callbacks[-1]}" ) if len(callbacks) > 2: second_to_last = callbacks[-2] assert second_to_last[0] < m, ( f"Rank {rank}: second-to-last callback should report < {m}, " f"got {second_to_last}" ) # All callbacks must have total == M for i, (_, total) in enumerate(callbacks): assert total == m, ( f"Rank {rank}: callback {i} has total={total}, expected {m}" ) # processed values must be non-decreasing processed_vals = [p for p, _ in callbacks] for i in range(1, len(processed_vals)): assert processed_vals[i] >= processed_vals[i - 1], ( f"Rank {rank}: callbacks not non-decreasing at index {i}: " f"{processed_vals}" ) # No duplicate consecutive callbacks (pipeline dummies must not emit callbacks) for i in range(1, len(callbacks)): assert callbacks[i] != callbacks[i - 1], ( f"Rank {rank}: duplicate consecutive callback at index {i}: " f"{callbacks[i]} (this suggests dummy iterations are emitting callbacks)" ) ================================================ FILE: src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py ================================================ import copy import gc import importlib import json import shutil import tempfile from dataclasses import dataclass from pathlib import Path from typing import Any, cast import mlx.core as mx import mlx.nn as nn import pytest from mlx.utils import tree_flatten, tree_unflatten from mlx_lm.tokenizer_utils import TokenizerWrapper from exo.shared.types.common import ModelId from exo.shared.types.mlx import Model from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams from exo.worker.engines.mlx.cache import KVPrefixCache from exo.worker.engines.mlx.generator.generate import mlx_generate from exo.worker.engines.mlx.utils_mlx import ( apply_chat_template, load_tokenizer_for_model_id, ) HF_CACHE = Path.home() / ".cache" / "huggingface" / "hub" # ── Config reduction ──────────────────────────────────────────────────────── # _REDUCE = { "num_hidden_layers": 4, "hidden_size": 256, "num_attention_heads": 4, "num_key_value_heads": 4, "intermediate_size": 512, "moe_intermediate_size": 128, "num_experts": 4, "num_experts_per_tok": 2, "n_routed_experts": 4, "num_local_experts": 4, "num_nextn_predict_layers": 0, "first_k_dense_replace": 0, "linear_num_key_heads": 2, "linear_num_value_heads": 2, "num_attention_groups": 4, } def _reduce_dict(cfg: dict[str, Any]) -> dict[str, Any]: result = dict(cfg) for key, val in _REDUCE.items(): if key in result: result[key] = val return result def _reduce_config(cfg: dict[str, Any]) -> dict[str, Any]: result = _reduce_dict(cfg) n_layers = cast(int, result.get("num_hidden_layers", 4)) if "text_config" in result and isinstance(result["text_config"], dict): result["text_config"] = _reduce_dict( cast(dict[str, Any], result["text_config"]) ) tc: dict[str, Any] = result["text_config"] if "num_nextn_predict_layers" in tc: tc["num_nextn_predict_layers"] = 0 if "layer_types" in result and isinstance(result["layer_types"], list): result["layer_types"] = result["layer_types"][:n_layers] if "attention_other_setting" in result and isinstance( result["attention_other_setting"], dict ): aos: dict[str, Any] = dict( cast(dict[str, Any], result["attention_other_setting"]) ) if "num_attention_heads" in aos: aos["num_attention_heads"] = result.get("num_attention_heads", 4) if "num_attention_groups" in aos: aos["num_attention_groups"] = result.get( "num_attention_groups", cast(int, aos["num_attention_groups"]) ) result["attention_other_setting"] = aos if "moe_layers_enum" in result and isinstance(result["moe_layers_enum"], str): indices = [int(x) for x in result["moe_layers_enum"].split(",") if x.strip()] valid = [i for i in indices if i < n_layers] result["moe_layers_enum"] = ",".join(str(i) for i in valid) if valid else "" return result # ── Helpers ───────────────────────────────────────────────────────────────── # def _find_snapshot(hub_name: str) -> Path | None: model_dir = HF_CACHE / f"models--mlx-community--{hub_name}" snaps = model_dir / "snapshots" if not snaps.exists(): return None children = sorted(snaps.iterdir()) return children[0] if children else None def _copy_tokenizer(src: Path, dst: Path) -> None: for f in src.iterdir(): name = f.name if ( "tokeniz" in name.lower() or "tiktoken" in name.lower() or name.startswith("vocab") or name.endswith(".jinja") or "tool_declaration" in name ) and f.is_file(): shutil.copy2(f, dst / name) def _build_model(module_name: str, cfg: dict[str, Any]) -> Model: mod = importlib.import_module(f"mlx_lm.models.{module_name}") args = mod.ModelArgs.from_dict(cfg) # pyright: ignore[reportAny] model: nn.Module = mod.Model(args) # pyright: ignore[reportAny] flat = cast(list[tuple[str, mx.array]], tree_flatten(model.parameters())) random_weights = [ (k, mx.random.normal(shape=v.shape, dtype=mx.float16)) for k, v in flat ] model.update(cast(dict[str, Any], tree_unflatten(random_weights))) mx.eval(model.parameters()) return cast(Model, model) def _collect_tokens( model: Model, tokenizer: TokenizerWrapper, task: TextGenerationTaskParams, prompt: str, kv_prefix_cache: KVPrefixCache | None, ) -> list[int]: tokens: list[int] = [] for resp in mlx_generate( model=model, tokenizer=tokenizer, task=task, prompt=prompt, kv_prefix_cache=kv_prefix_cache, group=None, ): tokens.append(resp.token) if resp.finish_reason is not None: break return tokens # ── Architecture definitions ──────────────────────────────────────────────── # @dataclass(frozen=True) class ArchSpec: name: str hub_name: str module: str tokenizer_hub: str | None = None # fallback for models without bundled tokenizer ARCHITECTURES: list[ArchSpec] = [ ArchSpec("llama", "Llama-3.2-1B-Instruct-4bit", "llama"), ArchSpec("glm_moe_dsa", "GLM-5-MXFP4-Q8", "glm_moe_dsa"), ArchSpec( "glm4_moe", "GLM-4.5-Air-8bit", "glm4_moe", tokenizer_hub="GLM-4.7-8bit-gs32" ), ArchSpec( "glm4_moe_lite", "GLM-4.7-Flash-8bit", "glm4_moe_lite", tokenizer_hub="GLM-4.7-8bit-gs32", ), ArchSpec("glm4_moe_47", "GLM-4.7-8bit-gs32", "glm4_moe"), ArchSpec("qwen3", "Qwen3-4B-Instruct-2507-4bit", "qwen3"), ArchSpec("qwen3_moe", "Qwen3-30B-A3B-4bit", "qwen3_moe"), ArchSpec("qwen3_next", "Qwen3-Next-80B-A3B-Thinking-4bit", "qwen3_next"), ArchSpec("minimax", "MiniMax-M2.1-3bit", "minimax"), ArchSpec("gpt_oss", "gpt-oss-20b-MXFP4-Q8", "gpt_oss"), ArchSpec("step3p5", "Step-3.5-Flash-4bit", "step3p5"), ArchSpec("kimi_k25", "Kimi-K2.5", "kimi_k25"), ArchSpec("qwen3_5", "Qwen3.5-2B-MLX-8bit", "qwen3_5"), ArchSpec("qwen3_5_moe", "Qwen3.5-35B-A3B-4bit", "qwen3_5_moe"), ] def _arch_available(spec: ArchSpec) -> bool: snap = _find_snapshot(spec.hub_name) if snap is None: return False if spec.tokenizer_hub is not None: return _find_snapshot(spec.tokenizer_hub) is not None return True def _make_task() -> TextGenerationTaskParams: return TextGenerationTaskParams( model=ModelId("test"), input=[ InputMessage( role="user", content="Use the calculator to compute 1847 * 263 + 5921", ) ], max_output_tokens=20, temperature=0.0, tools=[ { "type": "function", "function": { "name": "calculate", "description": "Evaluate a mathematical expression", "parameters": { "type": "object", "properties": {"expression": {"type": "string"}}, "required": ["expression"], }, }, } ], ) # ── Test class ────────────────────────────────────────────────────────────── # @pytest.mark.slow class TestPrefixCacheArchitectures: """Verify prefix cache produces identical output to fresh generation for every architecture.""" @pytest.fixture(autouse=True) def _cleanup(self): yield mx.clear_cache() gc.collect() @pytest.mark.parametrize( "spec", ARCHITECTURES, ids=[a.name for a in ARCHITECTURES], ) def test_prefix_cache_exact_hit(self, spec: ArchSpec) -> None: if not _arch_available(spec): pytest.skip(f"Model {spec.hub_name} not cached locally") snapshot = _find_snapshot(spec.hub_name) assert snapshot is not None tmpdir = Path(tempfile.mkdtemp(prefix=f"exo_test_{spec.name}_")) try: # Build reduced config with open(snapshot / "config.json") as f: cfg = cast(dict[str, Any], json.load(f)) reduced = _reduce_config(copy.deepcopy(cfg)) (tmpdir / "config.json").write_text(json.dumps(reduced)) # Copy tokenizer tok_src = snapshot if spec.tokenizer_hub is not None: alt = _find_snapshot(spec.tokenizer_hub) if alt is not None: tok_src = alt _copy_tokenizer(tok_src, tmpdir) # Load tokenizer and model model_id = ModelId(f"mlx-community/{spec.hub_name}") tokenizer = load_tokenizer_for_model_id(model_id, tmpdir) mx.random.seed(0) model = _build_model(spec.module, reduced) task = _make_task() prompt = apply_chat_template(tokenizer=tokenizer, task_params=task) # Run 1: fresh mx.random.seed(42) fresh = _collect_tokens(model, tokenizer, task, prompt, None) assert len(fresh) > 0, "Fresh generation produced no tokens" # Run 2: populate cache kv = KVPrefixCache(None) mx.random.seed(42) populate = _collect_tokens(model, tokenizer, task, prompt, kv) # Run 3: exact cache hit mx.random.seed(42) cached = _collect_tokens(model, tokenizer, task, prompt, kv) assert fresh == populate, ( f"Fresh vs populate mismatch: {fresh[:5]} vs {populate[:5]}" ) assert fresh == cached, ( f"Fresh vs cached mismatch: {fresh[:5]} vs {cached[:5]}" ) finally: shutil.rmtree(tmpdir, ignore_errors=True) ================================================ FILE: src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py ================================================ """ Unit tests for tokenizer loading and functionality across all supported models. This test downloads only tokenizer-related files (not full model weights) to verify that tokenizers can be loaded and used correctly for encoding/decoding. """ import asyncio import contextlib from pathlib import Path import pytest from exo.download.download_utils import ( download_file_with_retry, ensure_models_dir, fetch_file_list_with_cache, ) from exo.shared.models.model_cards import ModelCard, ModelId, get_model_cards from exo.worker.engines.mlx.utils_mlx import ( get_eos_token_ids_for_model, load_tokenizer_for_model_id, ) # Files needed for tokenizer functionality TOKENIZER_FILE_PATTERNS = [ "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.json", "vocab.txt", "merges.txt", "tiktoken.model", "added_tokens.json", "tokenizer.model", "tokenization_*.py", # Custom tokenizer implementations "tool_declaration_ts.py", # Dependency of tokenization_kimi.py ] def is_tokenizer_file(filename: str) -> bool: """Check if a file is needed for tokenizer functionality.""" for pattern in TOKENIZER_FILE_PATTERNS: if "*" in pattern: prefix = pattern.split("*")[0] suffix = pattern.split("*")[1] if filename.startswith(prefix) and filename.endswith(suffix): return True elif filename == pattern: return True return False async def download_tokenizer_files(model_id: ModelId) -> Path: """Download only the tokenizer-related files for a model.""" target_dir = await ensure_models_dir() / model_id.normalize() target_dir.mkdir(parents=True, exist_ok=True) file_list = await fetch_file_list_with_cache(model_id, "main", recursive=True) tokenizer_files = [f for f in file_list if is_tokenizer_file(f.path)] if not tokenizer_files: pytest.skip(f"No tokenizer files found for {model_id}") for file_entry in tokenizer_files: with contextlib.suppress(FileNotFoundError): await download_file_with_retry( model_id, "main", file_entry.path, target_dir ) return target_dir # Get a sample of models to test (one per family to keep tests fast) def get_test_models() -> list[ModelCard]: """Get a representative sample of models to test.""" # Pick one model from each family to test families: dict[str, ModelCard] = {} for card in asyncio.run(get_model_cards()): # Extract family name (e.g., "llama-3.1" from "llama-3.1-8b") parts = card.model_id.short().split("-") family = "-".join(parts[:2]) if len(parts) >= 2 else parts[0] if family not in families: families[family] = card return list(families.values()) TEST_MODELS: list[ModelCard] = get_test_models() pytestmark = pytest.mark.slow @pytest.fixture(scope="module") def event_loop(): """Create event loop for async tests.""" loop = asyncio.new_event_loop() yield loop loop.close() @pytest.mark.parametrize( "model_card", TEST_MODELS, ) @pytest.mark.asyncio async def test_tokenizer_encode_decode(model_card: ModelCard) -> None: """Test that tokenizer can encode and decode text correctly.""" model_id = model_card.model_id # Download tokenizer files model_path = await download_tokenizer_files(model_id) # Verify required files exist has_tokenizer = ( (model_path / "tokenizer.json").exists() or (model_path / "tokenizer_config.json").exists() or (model_path / "tiktoken.model").exists() or (model_path / "tokenizer.model").exists() ) if not has_tokenizer: pytest.skip(f"Required tokenizer files not found for {model_id}") # Load tokenizer tokenizer = load_tokenizer_for_model_id(model_id, model_path) # Test basic encoding test_text = "Hello, world!" encoded = tokenizer.encode(test_text) assert isinstance(encoded, list), f"encode() should return a list for {model_id}" assert len(encoded) > 0, f"encode() should return non-empty list for {model_id}" assert all(isinstance(t, int) for t in encoded), ( f"All tokens should be integers for {model_id}" ) # Test decoding decoded = tokenizer.decode(encoded) assert isinstance(decoded, str), f"decode() should return a string for {model_id}" assert test_text in decoded or decoded.strip() == test_text.strip(), ( f"decode(encode(x)) should preserve text for {model_id}: got {decoded!r}" ) # Test with longer text long_text = "The quick brown fox jumps over the lazy dog. " * 10 long_encoded = tokenizer.encode(long_text) assert len(long_encoded) > len(encoded), ( f"Longer text should produce more tokens for {model_id}" ) # Test empty string empty_encoded = tokenizer.encode("") assert isinstance(empty_encoded, list), ( f"encode('') should return a list for {model_id}" ) # Test special characters special_text = 'Hello!\n\tWorld? & "quotes"' special_encoded = tokenizer.encode(special_text) assert len(special_encoded) > 0, f"Special chars should encode for {model_id}" # Test unicode unicode_text = "Hello 世界 🌍" unicode_encoded = tokenizer.encode(unicode_text) assert len(unicode_encoded) > 0, f"Unicode should encode for {model_id}" @pytest.mark.parametrize( "model_card", TEST_MODELS, ) @pytest.mark.asyncio async def test_tokenizer_has_required_attributes(model_card: ModelCard) -> None: """Test that tokenizer has required attributes for inference.""" model_id = model_card.model_id model_path = await download_tokenizer_files(model_id) has_tokenizer = ( (model_path / "tokenizer.json").exists() or (model_path / "tokenizer_config.json").exists() or (model_path / "tiktoken.model").exists() or (model_path / "tokenizer.model").exists() ) if not has_tokenizer: pytest.skip(f"Required tokenizer files not found for {model_id}") tokenizer = load_tokenizer_for_model_id(model_id, model_path) eos_token_ids = get_eos_token_ids_for_model(model_id) # Check for vocabulary size empty_vocab: dict[str, int] = {} vocab_size: int = getattr(tokenizer, "vocab_size", None) or len( getattr(tokenizer, "get_vocab", lambda: empty_vocab)() ) assert vocab_size > 0, f"Tokenizer should have vocab_size > 0 for {model_id}" # Check for EOS token (either from tokenizer or explicitly provided) has_eos = ( eos_token_ids is not None or getattr(tokenizer, "eos_token_id", None) is not None or getattr(tokenizer, "eos_token", None) is not None ) assert has_eos, f"Tokenizer should have EOS token for {model_id}" @pytest.mark.parametrize( "model_card", TEST_MODELS, ) @pytest.mark.asyncio async def test_tokenizer_special_tokens(model_card: ModelCard) -> None: """Test that tokenizer can encode text containing special tokens. This is critical because the actual inference path uses prompts with special tokens from chat templates. If special tokens aren't handled correctly, encoding will fail. """ model_id = model_card.model_id model_path = await download_tokenizer_files(model_id) has_tokenizer = ( (model_path / "tokenizer.json").exists() or (model_path / "tokenizer_config.json").exists() or (model_path / "tiktoken.model").exists() or (model_path / "tokenizer.model").exists() ) assert has_tokenizer, f"Required tokenizer files not found for {model_id}" tokenizer = load_tokenizer_for_model_id(model_id, model_path) # Get special tokens from the tokenizer special_tokens: list[str] = [] # Try to get special tokens from various sources if hasattr(tokenizer, "all_special_tokens"): special_tokens.extend(tokenizer.all_special_tokens) elif hasattr(tokenizer, "_tokenizer") and hasattr( tokenizer._tokenizer, "all_special_tokens", ): special_tokens.extend(tokenizer._tokenizer.all_special_tokens) # Also check for common special token attributes for attr in [ "bos_token", "eos_token", "pad_token", "unk_token", "sep_token", "cls_token", ]: token = getattr(tokenizer, attr, None) if token is None and hasattr(tokenizer, "_tokenizer"): token = getattr(tokenizer._tokenizer, attr, None) if token and isinstance(token, str) and token not in special_tokens: special_tokens.append(token) # If we found special tokens, test encoding text that contains them if special_tokens: # Create text with special tokens interspersed test_with_special = f"{special_tokens[0]}Hello world" if len(special_tokens) > 1: test_with_special += f"{special_tokens[1]}" encoded = tokenizer.encode(test_with_special) assert isinstance(encoded, list), ( f"encode() with special tokens should return list for {model_id}" ) assert len(encoded) > 0, ( f"encode() with special tokens should return non-empty list for {model_id}" ) assert all(isinstance(t, int) for t in encoded), ( f"All tokens should be integers for {model_id}" ) # Verify we can decode decoded = tokenizer.decode(encoded) assert isinstance(decoded, str), f"decode() should return string for {model_id}" # Test with angle-bracket tokens (common format for special tokens) # These should not raise errors even if they're not actual special tokens angle_bracket_text = "<|test|>Hello<|end|>" encoded = tokenizer.encode(angle_bracket_text) assert isinstance(encoded, list), ( f"encode() with angle brackets should return list for {model_id}" ) assert len(encoded) > 0, ( f"encode() with angle brackets should be non-empty for {model_id}" ) # Specifically test Kimi tokenizer since it has special handling @pytest.mark.asyncio async def test_kimi_tokenizer_specifically(): """Test Kimi tokenizer with its specific patches and quirks.""" kimi_models = [ card for card in await get_model_cards() if "kimi" in card.model_id.lower() ] if not kimi_models: pytest.skip("No Kimi models found in MODEL_CARDS") model_card = kimi_models[0] model_id = model_card.model_id model_path = await download_tokenizer_files(model_id) # Ensure the custom tokenizer file exists if not (model_path / "tokenization_kimi.py").exists(): pytest.skip("tokenization_kimi.py not found") tokenizer = load_tokenizer_for_model_id(model_id, model_path) eos_token_ids = get_eos_token_ids_for_model(model_id) # Test encode/decode cycle test_text = "Hello, world!" encoded = tokenizer.encode(test_text) decoded = tokenizer.decode(encoded) assert len(encoded) > 0, "Kimi tokenizer should encode text" assert isinstance(decoded, str), "Kimi tokenizer should decode to string" # Test that the patched encode works (returns list of ints) assert all(isinstance(t, int) for t in encoded), "Tokens should be integers" # Test encoding text with special tokens (like from chat templates) # This is critical - the warmup inference uses prompts with special tokens special_token_text = "<|im_user|>user<|im_middle|>Hello<|im_end|><|im_assistant|>" special_encoded = tokenizer.encode(special_token_text) assert len(special_encoded) > 0, "Kimi tokenizer should handle special tokens" assert all(isinstance(t, int) for t in special_encoded), ( "Special token encoding should return integers" ) # Verify EOS token is set assert eos_token_ids == [163586], "Kimi EOS token should be [163586]" # Test GLM tokenizer since it also has special handling @pytest.mark.asyncio async def test_glm_tokenizer_specifically(): """Test GLM tokenizer with its specific EOS tokens.""" def contains(card: ModelCard, x: str): return x in card.model_id.lower() glm_model_cards = [ card for card in await get_model_cards() if contains(card, "glm") and not contains(card, "-5") and not contains(card, "4.7") ] if not glm_model_cards: pytest.skip("No GLM models found in MODEL_CARDS") model_card = glm_model_cards[0] model_id = model_card.model_id model_path = await download_tokenizer_files(model_id) has_tokenizer = (model_path / "tokenizer.json").exists() or ( model_path / "tokenizer_config.json" ).exists() if not has_tokenizer: pytest.skip("GLM tokenizer files not found") tokenizer = load_tokenizer_for_model_id(model_id, model_path) eos_token_ids = get_eos_token_ids_for_model(model_id) # Test encode/decode test_text = "Hello, world!" encoded = tokenizer.encode(test_text) decoded = tokenizer.decode(encoded) assert len(encoded) > 0, "GLM tokenizer should encode text" assert isinstance(decoded, str), "GLM tokenizer should decode to string" # Verify EOS tokens assert eos_token_ids == [ 151336, 151329, 151338, ], "GLM EOS tokens should be correct" ================================================ FILE: src/exo/worker/tests/unittests/test_plan/__init__.py ================================================ ================================================ FILE: src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py ================================================ import exo.worker.plan as plan_mod from exo.shared.types.common import NodeId from exo.shared.types.memory import Memory from exo.shared.types.tasks import LoadModel from exo.shared.types.worker.downloads import DownloadCompleted, DownloadProgress from exo.shared.types.worker.instances import BoundInstance from exo.shared.types.worker.runners import ( RunnerConnected, RunnerIdle, ) from exo.worker.tests.constants import ( INSTANCE_1_ID, MODEL_A_ID, NODE_A, NODE_B, RUNNER_1_ID, RUNNER_2_ID, ) from exo.worker.tests.unittests.conftest import ( FakeRunnerSupervisor, get_mlx_ring_instance, get_pipeline_shard_metadata, ) def test_plan_requests_download_when_waiting_and_shard_not_downloaded(): """ When a runner is waiting for a model and its shard is not in the local download_status map, plan() should emit DownloadModel. """ shard = get_pipeline_shard_metadata(model_id=MODEL_A_ID, device_rank=0) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID}, runner_to_shard={RUNNER_1_ID: shard}, ) bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) runner = FakeRunnerSupervisor(bound_instance=bound_instance, status=RunnerIdle()) runners = {RUNNER_1_ID: runner} instances = {INSTANCE_1_ID: instance} all_runners = {RUNNER_1_ID: RunnerIdle()} result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status={NODE_A: []}, instances=instances, all_runners=all_runners, tasks={}, ) assert isinstance(result, plan_mod.DownloadModel) assert result.instance_id == INSTANCE_1_ID assert result.shard_metadata == shard def test_plan_loads_model_when_all_shards_downloaded_and_waiting(): """ When all shards for an instance are DownloadCompleted (globally) and all runners are in waiting/loading/loaded states, plan() should emit LoadModel once. """ shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard2 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, runner_to_shard={RUNNER_1_ID: shard1, RUNNER_2_ID: shard2}, ) bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) local_runner = FakeRunnerSupervisor( bound_instance=bound_instance, status=RunnerConnected() ) runners = {RUNNER_1_ID: local_runner} instances = {INSTANCE_1_ID: instance} all_runners = { RUNNER_1_ID: RunnerConnected(), RUNNER_2_ID: RunnerConnected(), } global_download_status = { NODE_A: [ DownloadCompleted(shard_metadata=shard1, node_id=NODE_A, total=Memory()) ], NODE_B: [ DownloadCompleted(shard_metadata=shard2, node_id=NODE_B, total=Memory()) ], } result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status=global_download_status, instances=instances, all_runners=all_runners, tasks={}, ) assert isinstance(result, LoadModel) assert result.instance_id == INSTANCE_1_ID def test_plan_does_not_request_download_when_shard_already_downloaded(): """ If the local shard already has a DownloadCompleted entry, plan() should not re-emit DownloadModel while global state is still catching up. """ shard = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID}, runner_to_shard={RUNNER_1_ID: shard}, ) bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) runner = FakeRunnerSupervisor(bound_instance=bound_instance, status=RunnerIdle()) runners = {RUNNER_1_ID: runner} instances = {INSTANCE_1_ID: instance} all_runners = {RUNNER_1_ID: RunnerIdle()} # Global state shows shard is downloaded for NODE_A global_download_status: dict[NodeId, list[DownloadProgress]] = { NODE_A: [ DownloadCompleted(shard_metadata=shard, node_id=NODE_A, total=Memory()) ], NODE_B: [], } result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status=global_download_status, instances=instances, all_runners=all_runners, tasks={}, ) assert not isinstance(result, plan_mod.DownloadModel) def test_plan_does_not_load_model_until_all_shards_downloaded_globally(): """ LoadModel should not be emitted while some shards are still missing from the global_download_status. """ shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard2 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, runner_to_shard={RUNNER_1_ID: shard1, RUNNER_2_ID: shard2}, ) bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) local_runner = FakeRunnerSupervisor( bound_instance=bound_instance, status=RunnerConnected() ) runners = {RUNNER_1_ID: local_runner} instances = {INSTANCE_1_ID: instance} all_runners = { RUNNER_1_ID: RunnerConnected(), RUNNER_2_ID: RunnerConnected(), } global_download_status = { NODE_A: [ DownloadCompleted(shard_metadata=shard1, node_id=NODE_A, total=Memory()) ], NODE_B: [], # NODE_B has no downloads completed yet } result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status=global_download_status, instances=instances, all_runners=all_runners, tasks={}, ) assert result is None global_download_status = { NODE_A: [ DownloadCompleted(shard_metadata=shard1, node_id=NODE_A, total=Memory()) ], NODE_B: [ DownloadCompleted(shard_metadata=shard2, node_id=NODE_B, total=Memory()) ], # NODE_B has no downloads completed yet } result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status=global_download_status, instances=instances, all_runners=all_runners, tasks={}, ) assert result is not None ================================================ FILE: src/exo/worker/tests/unittests/test_plan/test_runner_lifecycle.py ================================================ from typing import Any import exo.worker.plan as plan_mod from exo.shared.types.tasks import Shutdown from exo.shared.types.worker.instances import BoundInstance, Instance, InstanceId from exo.shared.types.worker.runners import ( RunnerFailed, RunnerId, RunnerReady, RunnerStatus, ) from exo.worker.tests.constants import ( INSTANCE_1_ID, MODEL_A_ID, NODE_A, NODE_B, RUNNER_1_ID, RUNNER_2_ID, ) from exo.worker.tests.unittests.conftest import ( FakeRunnerSupervisor, get_mlx_ring_instance, get_pipeline_shard_metadata, ) def test_plan_kills_runner_when_instance_missing(): """ If a local runner's instance is no longer present in state, plan() should return a Shutdown for that runner. """ shard = get_pipeline_shard_metadata(model_id=MODEL_A_ID, device_rank=0) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID}, runner_to_shard={RUNNER_1_ID: shard}, ) bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) runner = FakeRunnerSupervisor(bound_instance=bound_instance, status=RunnerReady()) runners = {RUNNER_1_ID: runner} instances: dict[InstanceId, Instance] = {} all_runners = {RUNNER_1_ID: RunnerReady()} result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore[arg-type] global_download_status={NODE_A: []}, instances=instances, all_runners=all_runners, tasks={}, ) assert isinstance(result, Shutdown) assert result.instance_id == INSTANCE_1_ID assert result.runner_id == RUNNER_1_ID def test_plan_kills_runner_when_sibling_failed(): """ If a sibling runner in the same instance has failed, the local runner should be shut down. """ shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard2 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, runner_to_shard={RUNNER_1_ID: shard1, RUNNER_2_ID: shard2}, ) bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) runner = FakeRunnerSupervisor(bound_instance=bound_instance, status=RunnerReady()) runners = {RUNNER_1_ID: runner} instances = {INSTANCE_1_ID: instance} all_runners = { RUNNER_1_ID: RunnerReady(), RUNNER_2_ID: RunnerFailed(error_message="boom"), } result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore[arg-type] global_download_status={NODE_A: []}, instances=instances, all_runners=all_runners, tasks={}, ) assert isinstance(result, Shutdown) assert result.instance_id == INSTANCE_1_ID assert result.runner_id == RUNNER_1_ID def test_plan_creates_runner_when_missing_for_node(): """ If shard_assignments specify a runner for this node but we don't have a local supervisor yet, plan() should emit a CreateRunner. """ shard = get_pipeline_shard_metadata(model_id=MODEL_A_ID, device_rank=0) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID}, runner_to_shard={RUNNER_1_ID: shard}, ) runners: dict[Any, Any] = {} # nothing local yet instances = {INSTANCE_1_ID: instance} all_runners: dict[Any, Any] = {} result = plan_mod.plan( node_id=NODE_A, runners=runners, global_download_status={NODE_A: []}, instances=instances, all_runners=all_runners, tasks={}, ) # We patched plan_mod.CreateRunner → CreateRunner assert isinstance(result, plan_mod.CreateRunner) assert result.instance_id == INSTANCE_1_ID assert isinstance(result.bound_instance, BoundInstance) assert result.bound_instance.instance is instance assert result.bound_instance.bound_runner_id == RUNNER_1_ID def test_plan_does_not_create_runner_when_supervisor_already_present(): """ If we already have a local supervisor for the runner assigned to this node, plan() should not emit a CreateRunner again. """ shard = get_pipeline_shard_metadata(model_id=MODEL_A_ID, device_rank=0) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID}, runner_to_shard={RUNNER_1_ID: shard}, ) bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) runner = FakeRunnerSupervisor(bound_instance=bound_instance, status=RunnerReady()) runners = {RUNNER_1_ID: runner} instances = {INSTANCE_1_ID: instance} all_runners = {RUNNER_1_ID: RunnerReady()} result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore[arg-type] global_download_status={NODE_A: []}, instances=instances, all_runners=all_runners, tasks={}, ) assert result is None def test_plan_does_not_create_runner_for_unassigned_node(): """ If this node does not appear in shard_assignments.node_to_runner, plan() should not try to create a runner on this node. """ shard = get_pipeline_shard_metadata(model_id=MODEL_A_ID, device_rank=0) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_B: RUNNER_2_ID}, runner_to_shard={RUNNER_2_ID: shard}, ) runners: dict[RunnerId, FakeRunnerSupervisor] = {} # no local runners instances = {INSTANCE_1_ID: instance} all_runners: dict[RunnerId, RunnerStatus] = {} result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status={NODE_A: []}, instances=instances, all_runners=all_runners, tasks={}, ) assert result is None ================================================ FILE: src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py ================================================ from typing import cast import exo.worker.plan as plan_mod from exo.shared.types.tasks import Task, TaskId, TaskStatus, TextGeneration from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams from exo.shared.types.worker.instances import BoundInstance, InstanceId from exo.shared.types.worker.runners import ( RunnerIdle, RunnerReady, RunnerRunning, ) from exo.worker.tests.constants import ( COMMAND_1_ID, INSTANCE_1_ID, MODEL_A_ID, NODE_A, NODE_B, RUNNER_1_ID, RUNNER_2_ID, TASK_1_ID, ) from exo.worker.tests.unittests.conftest import ( FakeRunnerSupervisor, OtherTask, get_mlx_ring_instance, get_pipeline_shard_metadata, ) def test_plan_forwards_pending_chat_completion_when_runner_ready(): """ When there is a pending TextGeneration for the local instance and all runners are Ready/Running, plan() should forward that task. """ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1}, ) bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) local_runner = FakeRunnerSupervisor( bound_instance=bound_instance, status=RunnerReady() ) runners = {RUNNER_1_ID: local_runner} instances = {INSTANCE_1_ID: instance} all_runners = { RUNNER_1_ID: RunnerReady(), RUNNER_2_ID: RunnerReady(), } task = TextGeneration( task_id=TASK_1_ID, instance_id=INSTANCE_1_ID, task_status=TaskStatus.Pending, command_id=COMMAND_1_ID, task_params=TextGenerationTaskParams( model=MODEL_A_ID, input=[InputMessage(role="user", content="")] ), ) result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status={NODE_A: []}, instances=instances, all_runners=all_runners, tasks={TASK_1_ID: task}, ) assert result is task def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready(): """ Even with a pending TextGeneration, plan() should not forward it unless all runners for the instance are Ready/Running. """ shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard2 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, runner_to_shard={RUNNER_1_ID: shard1, RUNNER_2_ID: shard2}, ) bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) local_runner = FakeRunnerSupervisor( bound_instance=bound_instance, status=RunnerReady() ) runners = {RUNNER_1_ID: local_runner} instances = {INSTANCE_1_ID: instance} all_runners = { RUNNER_1_ID: RunnerReady(), RUNNER_2_ID: RunnerIdle(), } task = TextGeneration( task_id=TASK_1_ID, instance_id=INSTANCE_1_ID, task_status=TaskStatus.Pending, command_id=COMMAND_1_ID, task_params=TextGenerationTaskParams( model=MODEL_A_ID, input=[InputMessage(role="user", content="")] ), ) result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status={NODE_A: [], NODE_B: []}, instances=instances, all_runners=all_runners, tasks={TASK_1_ID: task}, ) assert result is None def test_plan_does_not_forward_tasks_for_other_instances(): """ plan() should ignore pending TextGeneration tasks whose instance_id does not match the local instance. """ shard = get_pipeline_shard_metadata(model_id=MODEL_A_ID, device_rank=0) local_instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID}, runner_to_shard={RUNNER_1_ID: shard}, ) bound_instance = BoundInstance( instance=local_instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) local_runner = FakeRunnerSupervisor( bound_instance=bound_instance, status=RunnerReady() ) runners = {RUNNER_1_ID: local_runner} instances = {INSTANCE_1_ID: local_instance} all_runners = {RUNNER_1_ID: RunnerReady()} other_instance_id = InstanceId("instance-2") foreign_task = TextGeneration( task_id=TaskId("other-task"), instance_id=other_instance_id, task_status=TaskStatus.Pending, command_id=COMMAND_1_ID, task_params=TextGenerationTaskParams( model=MODEL_A_ID, input=[InputMessage(role="user", content="")] ), ) result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status={NODE_A: []}, instances=instances, all_runners=all_runners, tasks={foreign_task.task_id: foreign_task}, ) assert result is None def test_plan_ignores_non_pending_or_non_chat_tasks(): """ _pending_tasks should not forward tasks that are either not TextGeneration or not in Pending/Running states. """ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1}, ) bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) local_runner = FakeRunnerSupervisor( bound_instance=bound_instance, status=RunnerReady() ) runners = {RUNNER_1_ID: local_runner} instances = {INSTANCE_1_ID: instance} all_runners = { RUNNER_1_ID: RunnerReady(), RUNNER_2_ID: RunnerReady(), } completed_task = TextGeneration( task_id=TASK_1_ID, instance_id=INSTANCE_1_ID, task_status=TaskStatus.Complete, command_id=COMMAND_1_ID, task_params=TextGenerationTaskParams( model=MODEL_A_ID, input=[InputMessage(role="user", content="")] ), ) other_task_id = TaskId("other-task") other_task = cast( Task, cast( object, OtherTask( task_id=other_task_id, instance_id=INSTANCE_1_ID, task_status=TaskStatus.Pending, ), ), ) result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status={NODE_A: [], NODE_B: []}, instances=instances, all_runners=all_runners, tasks={TASK_1_ID: completed_task, other_task_id: other_task}, ) assert result is None def test_plan_returns_none_when_nothing_to_do(): """ If there are healthy runners, no downloads needed, and no pending tasks, plan() should return None (steady state). """ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1}, ) bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) local_runner = FakeRunnerSupervisor( bound_instance=bound_instance, status=RunnerRunning() ) runners = {RUNNER_1_ID: local_runner} instances = {INSTANCE_1_ID: instance} all_runners = { RUNNER_1_ID: RunnerRunning(), RUNNER_2_ID: RunnerRunning(), } result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status={NODE_A: [], NODE_B: []}, instances=instances, all_runners=all_runners, tasks={}, ) assert result is None ================================================ FILE: src/exo/worker/tests/unittests/test_plan/test_warmup.py ================================================ import exo.worker.plan as plan_mod from exo.shared.types.tasks import StartWarmup from exo.shared.types.worker.instances import BoundInstance from exo.shared.types.worker.runners import ( RunnerIdle, RunnerLoaded, RunnerLoading, RunnerWarmingUp, ) from exo.worker.tests.constants import ( INSTANCE_1_ID, MODEL_A_ID, NODE_A, NODE_B, NODE_C, RUNNER_1_ID, RUNNER_2_ID, RUNNER_3_ID, ) from exo.worker.tests.unittests.conftest import ( FakeRunnerSupervisor, get_mlx_ring_instance, get_pipeline_shard_metadata, ) def test_plan_starts_warmup_for_accepting_rank_when_all_loaded_or_warming(): """ For non-zero device_rank shards, StartWarmup should be emitted when all shards in the instance are Loaded/WarmingUp. """ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=3) shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=3) shard2 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=2, world_size=3) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID, NODE_C: RUNNER_3_ID}, runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1, RUNNER_3_ID: shard2}, ) bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_2_ID, bound_node_id=NODE_B ) local_runner = FakeRunnerSupervisor( bound_instance=bound_instance, status=RunnerLoaded() ) runners = {RUNNER_2_ID: local_runner} instances = {INSTANCE_1_ID: instance} all_runners = { RUNNER_1_ID: RunnerLoaded(), RUNNER_2_ID: RunnerLoaded(), RUNNER_3_ID: RunnerWarmingUp(), } result = plan_mod.plan( node_id=NODE_B, runners=runners, # type: ignore global_download_status={NODE_A: []}, instances=instances, all_runners=all_runners, tasks={}, ) assert isinstance(result, StartWarmup) assert result.instance_id == INSTANCE_1_ID def test_plan_starts_warmup_for_rank_zero_after_others_warming(): """ For device_rank == 0, StartWarmup should only be emitted once all the other runners in the instance are already warming up. """ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1}, ) bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) local_runner = FakeRunnerSupervisor( bound_instance=bound_instance, status=RunnerLoaded() ) runners = {RUNNER_1_ID: local_runner} instances = {INSTANCE_1_ID: instance} all_runners = { RUNNER_1_ID: RunnerLoaded(), RUNNER_2_ID: RunnerWarmingUp(), } result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status={NODE_A: []}, instances=instances, all_runners=all_runners, tasks={}, ) assert isinstance(result, StartWarmup) assert result.instance_id == INSTANCE_1_ID def test_plan_does_not_start_warmup_for_non_zero_rank_until_all_loaded_or_warming(): """ Non-zero rank should not start warmup while any shard is not Loaded/WarmingUp. """ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1}, ) bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_2_ID, bound_node_id=NODE_B ) local_runner = FakeRunnerSupervisor( bound_instance=bound_instance, status=RunnerLoaded() ) runners = {RUNNER_2_ID: local_runner} instances = {INSTANCE_1_ID: instance} all_runners = { RUNNER_1_ID: RunnerIdle(), RUNNER_2_ID: RunnerLoaded(), } result = plan_mod.plan( node_id=NODE_B, runners=runners, # type: ignore global_download_status={NODE_A: [], NODE_B: []}, instances=instances, all_runners=all_runners, tasks={}, ) assert result is None def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming(): """ Rank-zero shard should not start warmup until all non-zero ranks are already WarmingUp. For accepting ranks (device_rank != 0), StartWarmup should be emitted when all shards in the instance are Loaded/WarmingUp. In a 2-node setup, rank 1 is the accepting rank. """ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1}, ) # Rank 1 is the accepting rank bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) local_runner = FakeRunnerSupervisor( bound_instance=bound_instance, status=RunnerLoaded() ) runners = {RUNNER_1_ID: local_runner} instances = {INSTANCE_1_ID: instance} all_runners = { RUNNER_1_ID: RunnerLoaded(), RUNNER_2_ID: RunnerLoaded(), } result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status={NODE_A: []}, instances=instances, all_runners=all_runners, tasks={}, ) assert result is None all_runners = { RUNNER_1_ID: RunnerLoaded(), RUNNER_2_ID: RunnerWarmingUp(), } result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status={NODE_A: []}, instances=instances, all_runners=all_runners, tasks={}, ) assert isinstance(result, StartWarmup) assert result.instance_id == INSTANCE_1_ID def test_plan_starts_warmup_for_connecting_rank_after_others_warming(): """ For connecting rank (device_rank == world_size - 1), StartWarmup should only be emitted once all the other runners are already warming up. In a 2-node setup, rank 1 is the connecting rank. """ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1}, ) # Rank 1 is the connecting rank bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_2_ID, bound_node_id=NODE_B ) local_runner = FakeRunnerSupervisor( bound_instance=bound_instance, status=RunnerLoaded() ) runners = {RUNNER_2_ID: local_runner} instances = {INSTANCE_1_ID: instance} all_runners = { RUNNER_1_ID: RunnerWarmingUp(), RUNNER_2_ID: RunnerLoaded(), } result = plan_mod.plan( node_id=NODE_B, runners=runners, # type: ignore global_download_status={NODE_B: []}, instances=instances, all_runners=all_runners, tasks={}, ) assert isinstance(result, StartWarmup) assert result.instance_id == INSTANCE_1_ID def test_plan_does_not_start_warmup_for_accepting_rank_until_all_loaded_or_warming(): """ Accepting rank should not start warmup while any shard is not Loaded/WarmingUp. In a 2-node setup, rank 0 is the accepting rank. """ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1}, ) # Rank 0 is the accepting rank bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) local_runner = FakeRunnerSupervisor( bound_instance=bound_instance, status=RunnerLoaded() ) runners = {RUNNER_1_ID: local_runner} instances = {INSTANCE_1_ID: instance} all_runners = { RUNNER_1_ID: RunnerLoaded(), RUNNER_2_ID: RunnerLoading(), } result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status={NODE_A: [], NODE_B: []}, instances=instances, all_runners=all_runners, tasks={}, ) assert result is None def test_plan_does_not_start_warmup_for_connecting_rank_until_others_warming(): """ Connecting rank (device_rank == 0) should not start warmup until all other ranks are already WarmingUp. """ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) instance = get_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1}, ) # Rank 1 is the connecting rank bound_instance = BoundInstance( instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A ) local_runner = FakeRunnerSupervisor( bound_instance=bound_instance, status=RunnerLoaded() ) runners = {RUNNER_1_ID: local_runner} instances = {INSTANCE_1_ID: instance} all_runners = { RUNNER_1_ID: RunnerLoaded(), RUNNER_2_ID: RunnerLoaded(), } result = plan_mod.plan( node_id=NODE_A, runners=runners, # type: ignore global_download_status={NODE_A: [], NODE_B: []}, instances=instances, all_runners=all_runners, tasks={}, ) assert result is None ================================================ FILE: src/exo/worker/tests/unittests/test_runner/__init__.py ================================================ ================================================ FILE: src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.py ================================================ import json from collections.abc import Generator from typing import Any from exo.shared.types.worker.runner_response import ( GenerationResponse, ToolCallResponse, ) from exo.worker.engines.mlx.dsml_encoding import ( ASSISTANT_TOKEN, BOS_TOKEN, DSML_TOKEN, EOS_TOKEN, THINKING_END, THINKING_START, TOOL_CALLS_END, TOOL_CALLS_START, USER_TOKEN, encode_messages, parse_dsml_output, ) from exo.worker.runner.llm_inference.model_output_parsers import parse_deepseek_v32 # ── Shared fixtures ────────────────────────────────────────────── _WEATHER_TOOLS: list[dict[str, Any]] = [ { "type": "function", "function": { "name": "get_weather", "description": "Get the current weather in a given city", "parameters": { "type": "object", "properties": { "city": {"type": "string", "description": "The city name"}, "units": { "type": "string", "enum": ["celsius", "fahrenheit"], "description": "Temperature units", }, }, "required": ["city"], }, }, }, { "type": "function", "function": { "name": "get_time", "description": "Get the current time in a timezone", "parameters": { "type": "object", "properties": { "timezone": {"type": "string"}, }, "required": ["timezone"], }, }, }, ] def _simulate_tokens( texts: list[str], finish_on_last: bool = True, ) -> Generator[GenerationResponse]: """Simulate a model producing tokens from a list of text strings.""" for i, text in enumerate(texts): is_last = i == len(texts) - 1 yield GenerationResponse( text=text, token=i, finish_reason="stop" if (is_last and finish_on_last) else None, usage=None, ) # ── Test: Standard text response (no tool calls) ──────────────── class TestE2EStandardResponse: """Model generates a plain text response — no tool calling involved.""" def test_plain_text_passthrough(self): """Simulate model producing: 'The weather in NYC is 72°F and sunny.'""" # Step 1: Encode the prompt (with tools available) messages: list[dict[str, Any]] = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What's the weather in NYC?"}, ] prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS) # Verify prompt structure assert BOS_TOKEN in prompt assert "## Tools" in prompt assert "get_weather" in prompt assert f"{USER_TOKEN}What's the weather in NYC?{ASSISTANT_TOKEN}" in prompt # Step 2: Simulate model response — plain text tokens (no DSML) model_tokens = [ "The weather", " in NYC", " is 72", "°F", " and sunny", ".", ] results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) # Step 3: Verify all tokens pass through as GenerationResponse gen_results = [r for r in results if isinstance(r, GenerationResponse)] tool_results = [r for r in results if isinstance(r, ToolCallResponse)] assert len(tool_results) == 0 assert len(gen_results) == 6 full_text = "".join(r.text for r in gen_results) assert full_text == "The weather in NYC is 72°F and sunny." assert gen_results[-1].finish_reason == "stop" # ── Test: Tool call response ───────────────────────────────────── class TestE2EToolCallResponse: """Model generates a DSML tool call — realistic token boundaries.""" def test_realistic_tool_call_tokens(self): """Simulate model generating a get_weather tool call with realistic token splits. Real models split DSML markers across tokens unpredictably. This simulates how DeepSeek V3.2 actually tokenizes DSML output. """ # Step 1: Encode prompt messages: list[dict[str, Any]] = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What's the weather in San Francisco?"}, ] prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS) assert "get_weather" in prompt # Step 2: Simulate realistic token-by-token model output # The model first produces some text, then a DSML tool call block model_tokens = [ "I'll check the weather for you.", "\n\n", f"<{DSML_TOKEN}", # marker split across tokens "function_calls>\n", f'<{DSML_TOKEN}invoke name="get_weather">\n', f'<{DSML_TOKEN}parameter name="city" string="true">', "San Francisco", f"\n", f'<{DSML_TOKEN}parameter name="units" string="false">', '"celsius"', f"\n", f"\n", f"", ] results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) # Step 3: Verify gen_results = [r for r in results if isinstance(r, GenerationResponse)] tool_results = [r for r in results if isinstance(r, ToolCallResponse)] # Should have text tokens before tool call + one ToolCallResponse assert len(tool_results) == 1 assert len(tool_results[0].tool_calls) == 1 tc = tool_results[0].tool_calls[0] assert tc.name == "get_weather" args = json.loads(tc.arguments) # pyright: ignore[reportAny] assert args["city"] == "San Francisco" assert args["units"] == "celsius" # The text before the tool call should still be yielded text_before = "".join(r.text for r in gen_results if not r.is_thinking) assert "check the weather" in text_before def test_multiple_tool_calls_in_one_block(self): """Model generates two tool calls in a single function_calls block.""" messages: list[dict[str, Any]] = [ {"role": "system", "content": "You are helpful."}, {"role": "user", "content": "Weather in NYC and time in EST?"}, ] prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS) assert "get_weather" in prompt assert "get_time" in prompt # Simulate model output with two invocations model_tokens = [ "Let me check both.\n\n", TOOL_CALLS_START, "\n", f'<{DSML_TOKEN}invoke name="get_weather">\n', f'<{DSML_TOKEN}parameter name="city" string="true">NYC\n', f"\n", f'<{DSML_TOKEN}invoke name="get_time">\n', f'<{DSML_TOKEN}parameter name="timezone" string="true">EST\n', f"\n", TOOL_CALLS_END, ] results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) tool_results = [r for r in results if isinstance(r, ToolCallResponse)] assert len(tool_results) == 1 assert len(tool_results[0].tool_calls) == 2 assert tool_results[0].tool_calls[0].name == "get_weather" assert tool_results[0].tool_calls[1].name == "get_time" args0 = json.loads(tool_results[0].tool_calls[0].arguments) # pyright: ignore[reportAny] args1 = json.loads(tool_results[0].tool_calls[1].arguments) # pyright: ignore[reportAny] assert args0 == {"city": "NYC"} assert args1 == {"timezone": "EST"} # ── Test: Multi-turn tool use flow ─────────────────────────────── class TestE2EMultiTurnToolUse: """Full multi-turn: user asks → model calls tool → tool result → model answers.""" def test_encode_multi_turn_with_tool_results(self): """Verify the prompt for turn 2 (after tool results) is correctly encoded.""" # Turn 1: user asks, model calls tool # Turn 2: tool result provided, model answers messages: list[dict[str, Any]] = [ {"role": "system", "content": "You are a weather assistant."}, {"role": "user", "content": "What's the weather in NYC?"}, { "role": "assistant", "content": "", "tool_calls": [ { "type": "function", "function": { "name": "get_weather", "arguments": '{"city": "NYC"}', }, } ], }, {"role": "tool", "content": '{"temperature": 72, "condition": "sunny"}'}, ] prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS) # Verify multi-turn structure assert BOS_TOKEN in prompt assert "You are a weather assistant." in prompt assert "## Tools" in prompt # The assistant's tool call should be encoded as DSML assert TOOL_CALLS_START in prompt assert f'<{DSML_TOKEN}invoke name="get_weather">' in prompt assert EOS_TOKEN in prompt # The tool result should be wrapped in function_results assert "" in prompt assert "" in prompt assert "72" in prompt assert "" in prompt # Now simulate model answering after seeing the tool result model_tokens = [ "The current", " weather in NYC", " is 72°F", " and sunny.", ] results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) gen_results = [r for r in results if isinstance(r, GenerationResponse)] tool_results = [r for r in results if isinstance(r, ToolCallResponse)] assert len(tool_results) == 0 full_text = "".join(r.text for r in gen_results) assert full_text == "The current weather in NYC is 72°F and sunny." def test_multi_tool_results_encoding(self): """Verify encoding when model called two tools and both return results.""" messages: list[dict[str, Any]] = [ {"role": "user", "content": "Weather and time?"}, { "role": "assistant", "content": "", "tool_calls": [ { "type": "function", "function": { "name": "get_weather", "arguments": '{"city": "LA"}', }, }, { "type": "function", "function": { "name": "get_time", "arguments": '{"timezone": "PST"}', }, }, ], }, {"role": "tool", "content": "85F, clear skies"}, {"role": "tool", "content": "3:42 PM PST"}, ] prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS) # Should have one function_results block with two results assert prompt.count("") == 1 assert prompt.count("") == 1 assert "85F, clear skies" in prompt assert "3:42 PM PST" in prompt # ── Test: Thinking + tool call ─────────────────────────────────── class TestE2EThinkingAndToolCall: """Model uses thinking mode, reasons, then makes a tool call.""" def test_thinking_then_tool_call(self): """Model thinks first, then produces a DSML tool call block.""" messages: list[dict[str, Any]] = [ {"role": "user", "content": "What's the weather?"}, ] prompt = encode_messages( messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" ) # Thinking mode: prompt should end with assert prompt.endswith(THINKING_START) # Simulate: model outputs , thinks, closes thinking, then tool call. # In the full pipeline, parse_thinking_models handles the case where # is in the prompt. Here we test parse_deepseek_v32 directly, # which detects / markers in the stream. model_tokens = [ THINKING_START, "The user wants weather", " information. I should use", " the get_weather tool.", THINKING_END, "\n\n", TOOL_CALLS_START, "\n", f'<{DSML_TOKEN}invoke name="get_weather">\n', f'<{DSML_TOKEN}parameter name="city" string="true">', "San Francisco", f"\n", f"\n", TOOL_CALLS_END, ] results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) gen_results = [r for r in results if isinstance(r, GenerationResponse)] tool_results = [r for r in results if isinstance(r, ToolCallResponse)] # Should have thinking tokens + tool call thinking_results = [r for r in gen_results if r.is_thinking] assert len(thinking_results) >= 1 thinking_text = "".join(r.text for r in thinking_results) assert "get_weather tool" in thinking_text assert len(tool_results) == 1 assert tool_results[0].tool_calls[0].name == "get_weather" args = json.loads(tool_results[0].tool_calls[0].arguments) # pyright: ignore[reportAny] assert args["city"] == "San Francisco" def test_thinking_prompt_encoding(self): """Verify thinking mode affects prompt encoding correctly.""" messages: list[dict[str, Any]] = [ {"role": "system", "content": "Be thorough."}, {"role": "user", "content": "What's the weather?"}, ] # With thinking enabled prompt_think = encode_messages( messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" ) assert prompt_think.endswith(THINKING_START) # With thinking disabled prompt_no_think = encode_messages( messages, tools=_WEATHER_TOOLS, thinking_mode="chat" ) assert prompt_no_think.endswith(THINKING_END) # Both should have the same tool definitions assert "get_weather" in prompt_think assert "get_weather" in prompt_no_think # ── Test: Round-trip encode → parse ────────────────────────────── class TestE2ERoundTrip: """Verify that DSML we encode can be parsed back correctly.""" def test_encoded_tool_call_is_parseable(self): """Encode an assistant tool call message, then parse the DSML output.""" messages: list[dict[str, Any]] = [ {"role": "user", "content": "Weather?"}, { "role": "assistant", "content": "", "tool_calls": [ { "type": "function", "function": { "name": "get_weather", "arguments": '{"city": "Tokyo", "units": "celsius"}', }, } ], }, ] prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS) # Extract the DSML function_calls block from the prompt start = prompt.index(TOOL_CALLS_START) end = prompt.index(TOOL_CALLS_END) + len(TOOL_CALLS_END) dsml_block = prompt[start:end] # Parse it back parsed = parse_dsml_output(dsml_block) assert parsed is not None assert len(parsed) == 1 assert parsed[0].name == "get_weather" args = json.loads(parsed[0].arguments) # pyright: ignore[reportAny] assert args["city"] == "Tokyo" assert args["units"] == "celsius" def test_encoded_multi_tool_call_round_trips(self): """Encode multiple tool calls, verify they parse back correctly.""" messages: list[dict[str, Any]] = [ {"role": "user", "content": "Both please"}, { "role": "assistant", "content": "", "tool_calls": [ { "type": "function", "function": { "name": "get_weather", "arguments": '{"city": "Paris"}', }, }, { "type": "function", "function": { "name": "get_time", "arguments": '{"timezone": "CET"}', }, }, ], }, ] prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS) start = prompt.index(TOOL_CALLS_START) end = prompt.index(TOOL_CALLS_END) + len(TOOL_CALLS_END) dsml_block = prompt[start:end] parsed = parse_dsml_output(dsml_block) assert parsed is not None assert len(parsed) == 2 assert parsed[0].name == "get_weather" assert parsed[1].name == "get_time" assert json.loads(parsed[0].arguments) == {"city": "Paris"} assert json.loads(parsed[1].arguments) == {"timezone": "CET"} # ── Test: Edge cases with realistic token boundaries ───────────── class TestE2EEdgeCases: """Edge cases that occur in real model inference.""" def test_dsml_marker_split_at_fullwidth_pipe(self): """The fullwidth pipe character ｜ might be its own token.""" # This is a realistic tokenization: the DSML marker is split at the ｜ chars model_tokens = [ "Let me help.\n\n", "<\uff5c", # start of ｜DSML｜ "DSML\uff5c", # rest of DSML token "function_calls>\n", f'<{DSML_TOKEN}invoke name="get_weather">\n', f'<{DSML_TOKEN}parameter name="city" string="true">NYC\n', f"\n", TOOL_CALLS_END, ] results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) tool_results = [r for r in results if isinstance(r, ToolCallResponse)] assert len(tool_results) == 1 assert tool_results[0].tool_calls[0].name == "get_weather" def test_tool_call_with_nested_json_object(self): """Model passes a complex JSON object as a non-string parameter.""" dsml_block = ( f"{TOOL_CALLS_START}\n" f'<{DSML_TOKEN}invoke name="create_event">\n' f'<{DSML_TOKEN}parameter name="title" string="true">Team Standup\n' f'<{DSML_TOKEN}parameter name="config" string="false">' f'{{"recurring": true, "days": ["mon", "wed", "fri"], "time": "09:00"}}' f"\n" f"\n" f"{TOOL_CALLS_END}" ) # Feed as single token (model might produce it all at once after prefill) results = list(parse_deepseek_v32(_simulate_tokens([dsml_block]))) tool_results = [r for r in results if isinstance(r, ToolCallResponse)] assert len(tool_results) == 1 tc = tool_results[0].tool_calls[0] assert tc.name == "create_event" args = json.loads(tc.arguments) # pyright: ignore[reportAny] assert args["title"] == "Team Standup" assert args["config"]["recurring"] is True assert args["config"]["days"] == ["mon", "wed", "fri"] def test_text_with_angle_brackets_not_mistaken_for_dsml(self): """Angle brackets in normal text should not trigger DSML buffering.""" model_tokens = [ "The formula is ", "", " where x > 0", " and y < 100.", ] results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) gen_results = [r for r in results if isinstance(r, GenerationResponse)] tool_results = [r for r in results if isinstance(r, ToolCallResponse)] assert len(tool_results) == 0 full_text = "".join(r.text for r in gen_results) assert "formula" in full_text assert "" in full_text def test_empty_model_response(self): """Model produces only EOS (empty response).""" model_tokens = [""] results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) gen_results = [r for r in results if isinstance(r, GenerationResponse)] assert len(gen_results) == 1 assert gen_results[0].text == "" assert gen_results[0].finish_reason == "stop" # ── Test: Full EPDP spec round-trip ────────────────────────────── class TestE2EFullRoundTrip: """Full round-trip matching the vLLM EPDP spec. Simulates the complete multi-turn flow: Turn 1: user asks → think → tool call → tool result → think → answer Turn 2: user asks again → old reasoning stripped → think → answer """ def test_single_tool_full_flow_with_thinking(self): """Complete flow: user → think → tool call → tool result → think → answer. This is the core EPDP flow from the vLLM spec. """ # ── Turn 1.1: User asks, encode prompt ── messages: list[dict[str, Any]] = [ {"role": "system", "content": "You are a weather assistant."}, {"role": "user", "content": "How's the weather in Hangzhou?"}, ] prompt_1 = encode_messages( messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" ) assert prompt_1.endswith(THINKING_START) assert "## Tools" in prompt_1 assert "get_weather" in prompt_1 # ── Turn 1.1: Model thinks, then calls tool ── model_tokens_1 = [ THINKING_START, "The user wants to know the weather in Hangzhou.", " I need to use the get_weather tool.", THINKING_END, "\n\n", TOOL_CALLS_START, "\n", f'<{DSML_TOKEN}invoke name="get_weather">\n', f'<{DSML_TOKEN}parameter name="city" string="true">Hangzhou\n', f"\n", TOOL_CALLS_END, ] results_1 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_1))) # Verify: thinking tokens + tool call gen_1 = [r for r in results_1 if isinstance(r, GenerationResponse)] tool_1 = [r for r in results_1 if isinstance(r, ToolCallResponse)] thinking_1 = [r for r in gen_1 if r.is_thinking] assert len(thinking_1) >= 1 assert "get_weather tool" in "".join(r.text for r in thinking_1) assert len(tool_1) == 1 assert tool_1[0].tool_calls[0].name == "get_weather" tc_args = json.loads(tool_1[0].tool_calls[0].arguments) # pyright: ignore[reportAny] assert tc_args == {"city": "Hangzhou"} # ── Turn 1.2: Add assistant response + tool result to messages ── messages.append( { "role": "assistant", "content": "", "reasoning_content": "The user wants to know the weather in Hangzhou. I need to use the get_weather tool.", "tool_calls": [ { "type": "function", "function": { "name": "get_weather", "arguments": '{"city": "Hangzhou"}', }, } ], } ) messages.append( { "role": "tool", "content": '{"temperature": "7~13°C", "condition": "Cloudy"}', } ) # Encode prompt for turn 1.2 prompt_2 = encode_messages( messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" ) # Verify: prompt has the full conversation structure assert TOOL_CALLS_START in prompt_2 # assistant's encoded tool call assert EOS_TOKEN in prompt_2 # assistant turn ends with EOS assert "" in prompt_2 assert "" in prompt_2 assert "Cloudy" in prompt_2 assert "" in prompt_2 # After tool results with thinking enabled → appended assert prompt_2.endswith(THINKING_START) # The assistant's reasoning_content should appear (it's after last_user_idx) assert "get_weather tool" in prompt_2 # ── Turn 1.2: Model thinks about results, then answers ── model_tokens_2 = [ THINKING_START, "The weather in Hangzhou is Cloudy, 7~13°C.", " I'll tell the user.", THINKING_END, "The weather in Hangzhou is currently cloudy with temperatures between 7°C and 13°C.", ] results_2 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_2))) gen_2 = [r for r in results_2 if isinstance(r, GenerationResponse)] tool_2 = [r for r in results_2 if isinstance(r, ToolCallResponse)] thinking_2 = [r for r in gen_2 if r.is_thinking] non_thinking_2 = [r for r in gen_2 if not r.is_thinking] assert len(tool_2) == 0 # No more tool calls assert len(thinking_2) >= 1 assert "Cloudy" in "".join(r.text for r in thinking_2) assert len(non_thinking_2) >= 1 final_text = "".join(r.text for r in non_thinking_2) assert "7°C" in final_text assert "13°C" in final_text def test_multi_tool_full_flow(self): """Flow with two tools: user → think → 2 tool calls → 2 results → think → answer.""" # ── Initial prompt ── messages: list[dict[str, Any]] = [ {"role": "system", "content": "You help with weather and time."}, {"role": "user", "content": "Weather in NYC and time in EST?"}, ] prompt_1 = encode_messages( messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" ) assert prompt_1.endswith(THINKING_START) # ── Model thinks, calls both tools ── model_tokens_1 = [ THINKING_START, "Two requests: weather and time. I'll call both.", THINKING_END, "\n\n", TOOL_CALLS_START, "\n", f'<{DSML_TOKEN}invoke name="get_weather">\n', f'<{DSML_TOKEN}parameter name="city" string="true">NYC\n', f"\n", f'<{DSML_TOKEN}invoke name="get_time">\n', f'<{DSML_TOKEN}parameter name="timezone" string="true">EST\n', f"\n", TOOL_CALLS_END, ] results_1 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_1))) tool_1 = [r for r in results_1 if isinstance(r, ToolCallResponse)] assert len(tool_1) == 1 assert len(tool_1[0].tool_calls) == 2 assert tool_1[0].tool_calls[0].name == "get_weather" assert tool_1[0].tool_calls[1].name == "get_time" # ── Add assistant + both tool results ── messages.append( { "role": "assistant", "content": "", "reasoning_content": "Two requests: weather and time. I'll call both.", "tool_calls": [ { "type": "function", "function": { "name": "get_weather", "arguments": '{"city": "NYC"}', }, }, { "type": "function", "function": { "name": "get_time", "arguments": '{"timezone": "EST"}', }, }, ], } ) messages.append({"role": "tool", "content": "72°F, sunny"}) messages.append({"role": "tool", "content": "2:30 PM EST"}) prompt_2 = encode_messages( messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" ) # Verify multi-tool result encoding # Count is 2: 1 in _TOOLS_SYSTEM_TEMPLATE example + 1 in conversation assert prompt_2.count("") == 2 assert prompt_2.count("") == 2 assert "72°F, sunny" in prompt_2 assert "2:30 PM EST" in prompt_2 assert prompt_2.endswith(THINKING_START) # ── Model thinks about results, answers ── model_tokens_2 = [ THINKING_START, "Got both results. Weather is 72°F sunny, time is 2:30 PM.", THINKING_END, "In NYC it's currently 72°F and sunny. The time in EST is 2:30 PM.", ] results_2 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_2))) tool_2 = [r for r in results_2 if isinstance(r, ToolCallResponse)] gen_2 = [r for r in results_2 if isinstance(r, GenerationResponse)] non_thinking_2 = [r for r in gen_2 if not r.is_thinking] assert len(tool_2) == 0 final_text = "".join(r.text for r in non_thinking_2) assert "72°F" in final_text assert "2:30 PM" in final_text def test_two_user_turns_reasoning_stripped(self): """Turn 2: old reasoning_content is stripped from history. Per the vLLM spec, clear_reasoning_content is called between user turns to save bandwidth. Our _drop_old_thinking handles this. """ # Full turn 1 conversation (already completed) messages: list[dict[str, Any]] = [ {"role": "system", "content": "You are helpful."}, {"role": "user", "content": "Weather in Hangzhou?"}, { "role": "assistant", "content": "", "reasoning_content": "I need to call get_weather for Hangzhou.", "tool_calls": [ { "type": "function", "function": { "name": "get_weather", "arguments": '{"city": "Hangzhou"}', }, } ], }, {"role": "tool", "content": "Cloudy 7~13°C"}, { "role": "assistant", "content": "The weather in Hangzhou is cloudy, 7-13°C.", "reasoning_content": "The tool returned cloudy weather. I'll summarize.", }, # Turn 2: user asks again {"role": "user", "content": "What about Beijing?"}, ] prompt = encode_messages( messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" ) # Old reasoning_content from turn 1 assistants should be STRIPPED # (they're before the last user message at index 5) assert "I need to call get_weather" not in prompt assert "tool returned cloudy" not in prompt # But the assistant's content and tool calls should still be there assert "cloudy, 7-13°C" in prompt assert TOOL_CALLS_START in prompt # Prompt ends with for the new turn assert prompt.endswith(THINKING_START) # ── Turn 2: Model thinks, calls tool for Beijing ── model_tokens = [ THINKING_START, "Now the user wants Beijing weather.", THINKING_END, "\n\n", TOOL_CALLS_START, "\n", f'<{DSML_TOKEN}invoke name="get_weather">\n', f'<{DSML_TOKEN}parameter name="city" string="true">Beijing\n', f"\n", TOOL_CALLS_END, ] results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) tool_results = [r for r in results if isinstance(r, ToolCallResponse)] assert len(tool_results) == 1 assert tool_results[0].tool_calls[0].name == "get_weather" args = json.loads(tool_results[0].tool_calls[0].arguments) # pyright: ignore[reportAny] assert args == {"city": "Beijing"} def test_chained_tool_calls_loop(self): """Model calls tool, gets result, calls another tool, gets result, answers. This simulates the inner while loop from the vLLM spec where the model may need multiple sub-turns of tool calling before it has enough info. """ # ── Sub-turn 1: user asks, model calls get_time ── messages: list[dict[str, Any]] = [ {"role": "system", "content": "You are helpful."}, {"role": "user", "content": "What's the weather in Hangzhou tomorrow?"}, ] prompt_1 = encode_messages( messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" ) assert prompt_1.endswith(THINKING_START) # Model first calls get_time to figure out the date model_tokens_1 = [ THINKING_START, "I need the current date first to calculate tomorrow.", THINKING_END, "\n\n", TOOL_CALLS_START, "\n", f'<{DSML_TOKEN}invoke name="get_time">\n', f'<{DSML_TOKEN}parameter name="timezone" string="true">Asia/Shanghai\n', f"\n", TOOL_CALLS_END, ] results_1 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_1))) tool_1 = [r for r in results_1 if isinstance(r, ToolCallResponse)] assert len(tool_1) == 1 assert tool_1[0].tool_calls[0].name == "get_time" # ── Sub-turn 2: add tool result, model calls get_weather ── messages.append( { "role": "assistant", "content": "", "reasoning_content": "I need the current date first to calculate tomorrow.", "tool_calls": [ { "type": "function", "function": { "name": "get_time", "arguments": '{"timezone": "Asia/Shanghai"}', }, } ], } ) messages.append({"role": "tool", "content": "2025-12-01 14:30 CST"}) prompt_2 = encode_messages( messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" ) assert "2025-12-01 14:30 CST" in prompt_2 assert prompt_2.endswith(THINKING_START) # Model now knows the date, calls get_weather model_tokens_2 = [ THINKING_START, "Today is 2025-12-01, so tomorrow is 2025-12-02.", " Now I can check weather for Hangzhou.", THINKING_END, "\n\n", TOOL_CALLS_START, "\n", f'<{DSML_TOKEN}invoke name="get_weather">\n', f'<{DSML_TOKEN}parameter name="city" string="true">Hangzhou\n', f"\n", TOOL_CALLS_END, ] results_2 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_2))) tool_2 = [r for r in results_2 if isinstance(r, ToolCallResponse)] assert len(tool_2) == 1 assert tool_2[0].tool_calls[0].name == "get_weather" # ── Sub-turn 3: add weather result, model answers ── messages.append( { "role": "assistant", "content": "", "reasoning_content": "Today is 2025-12-01, so tomorrow is 2025-12-02. Now I can check weather for Hangzhou.", "tool_calls": [ { "type": "function", "function": { "name": "get_weather", "arguments": '{"city": "Hangzhou"}', }, } ], } ) messages.append({"role": "tool", "content": "Sunny, 5~12°C"}) prompt_3 = encode_messages( messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" ) # Should have both function_results blocks (one per tool round) # Count is 3: 1 in _TOOLS_SYSTEM_TEMPLATE example + 2 in conversation assert prompt_3.count("") == 3 assert prompt_3.count("") == 3 assert "2025-12-01 14:30 CST" in prompt_3 assert "Sunny, 5~12°C" in prompt_3 assert prompt_3.endswith(THINKING_START) # Model finally answers model_tokens_3 = [ THINKING_START, "I have the weather for tomorrow in Hangzhou.", THINKING_END, "Tomorrow in Hangzhou will be sunny with temperatures between 5°C and 12°C.", ] results_3 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_3))) tool_3 = [r for r in results_3 if isinstance(r, ToolCallResponse)] gen_3 = [r for r in results_3 if isinstance(r, GenerationResponse)] non_thinking_3 = [r for r in gen_3 if not r.is_thinking] assert len(tool_3) == 0 # No more tool calls — loop ends final_text = "".join(r.text for r in non_thinking_3) assert "sunny" in final_text.lower() assert "5°C" in final_text assert "12°C" in final_text ================================================ FILE: src/exo/worker/tests/unittests/test_runner/test_event_ordering.py ================================================ # Check tasks are complete before runner is ever ready. import unittest.mock from collections.abc import Iterable from typing import Callable import mlx.core as mx import pytest import exo.worker.runner.llm_inference.batch_generator as mlx_batch_generator import exo.worker.runner.llm_inference.model_output_parsers as mlx_model_output_parsers import exo.worker.runner.llm_inference.runner as mlx_runner from exo.shared.types.chunks import TokenChunk from exo.shared.types.events import ( ChunkGenerated, Event, RunnerStatusUpdated, TaskAcknowledged, TaskStatusUpdated, ) from exo.shared.types.tasks import ( ConnectToGroup, LoadModel, Shutdown, StartWarmup, Task, TaskId, TaskStatus, TextGeneration, ) from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams from exo.shared.types.worker.runner_response import GenerationResponse from exo.shared.types.worker.runners import ( RunnerConnected, RunnerConnecting, RunnerIdle, RunnerLoaded, RunnerLoading, RunnerReady, RunnerRunning, RunnerShutdown, RunnerShuttingDown, RunnerWarmingUp, ) from exo.utils.channels import mp_channel from ...constants import ( CHAT_COMPLETION_TASK_ID, COMMAND_1_ID, INITIALIZATION_TASK_ID, INSTANCE_1_ID, LOAD_TASK_ID, MODEL_A_ID, NODE_A, RUNNER_1_ID, SHUTDOWN_TASK_ID, WARMUP_TASK_ID, ) from ..conftest import get_bound_mlx_ring_instance def make_nothin[T, U, V](res: T) -> Callable[[], T]: def nothin(*_1: U, **_2: V) -> T: return res return nothin nothin = make_nothin(None) INIT_TASK = ConnectToGroup( task_id=INITIALIZATION_TASK_ID, instance_id=INSTANCE_1_ID, ) LOAD_TASK = LoadModel( task_id=LOAD_TASK_ID, instance_id=INSTANCE_1_ID, ) WARMUP_TASK = StartWarmup( task_id=WARMUP_TASK_ID, instance_id=INSTANCE_1_ID, ) SHUTDOWN_TASK = Shutdown( task_id=SHUTDOWN_TASK_ID, instance_id=INSTANCE_1_ID, runner_id=RUNNER_1_ID, ) CHAT_PARAMS = TextGenerationTaskParams( model=MODEL_A_ID, input=[InputMessage(role="user", content="hello")], stream=True, max_output_tokens=4, temperature=0.0, ) CHAT_TASK = TextGeneration( task_id=CHAT_COMPLETION_TASK_ID, command_id=COMMAND_1_ID, task_params=CHAT_PARAMS, instance_id=INSTANCE_1_ID, ) def assert_events_equal(test_events: Iterable[Event], true_events: Iterable[Event]): for test_event, true_event in zip(test_events, true_events, strict=True): test_event.event_id = true_event.event_id assert test_event == true_event, f"{test_event} != {true_event}" @pytest.fixture def patch_out_mlx(monkeypatch: pytest.MonkeyPatch): # initialize_mlx returns a mock group monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MockGroup())) monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, MockTokenizer))) monkeypatch.setattr(mlx_batch_generator, "warmup_inference", make_nothin(1)) monkeypatch.setattr(mlx_batch_generator, "_check_for_debug_prompts", nothin) monkeypatch.setattr(mlx_batch_generator, "mx_any", make_nothin(False)) def fake_all_gather( tasks: list[TextGeneration], group: object ) -> tuple[list[TextGeneration], list[TextGeneration]]: return (tasks, []) monkeypatch.setattr(mlx_batch_generator, "mx_all_gather_tasks", fake_all_gather) # Mock apply_chat_template since we're using a fake tokenizer (integer 1). # Returns a prompt without thinking tag so detect_thinking_prompt_suffix returns None. monkeypatch.setattr( mlx_batch_generator, "apply_chat_template", make_nothin("test prompt") ) monkeypatch.setattr( mlx_model_output_parsers, "detect_thinking_prompt_suffix", make_nothin(False) ) monkeypatch.setattr(mlx_batch_generator, "ExoBatchGenerator", FakeExoBatchGenerator) class FakeExoBatchGenerator: def __init__(self, *_args: object, **_kwargs: object) -> None: self._uid_counter = 0 self._pending: dict[int, GenerationResponse] = {} @property def has_work(self) -> bool: return bool(self._pending) def submit( self, task_params: object = None, prompt: object = None, on_prefill_progress: object = None, distributed_prompt_progress_callback: object = None, on_generation_token: object = None, ) -> int: uid = self._uid_counter self._uid_counter += 1 self._pending[uid] = GenerationResponse( text="hi", token=0, finish_reason="stop", usage=None, ) return uid def step(self) -> list[tuple[int, GenerationResponse]]: results = list(self._pending.items()) self._pending.clear() return results def cancel(self, uids: list[int]) -> None: for uid in uids: self._pending.pop(uid, None) def close(self) -> None: pass # Use a fake event_sender to remove test flakiness. class EventCollector: def __init__(self, on_event: Callable[[Event], None] | None = None) -> None: self.events: list[Event] = [] self._on_event = on_event def send(self, event: Event) -> None: self.events.append(event) if self._on_event: self._on_event(event) def close(self) -> None: pass def join(self) -> None: pass class MockTokenizer: tool_parser = None tool_call_start = None tool_call_end = None has_tool_calling = False has_thinking = False think_start = None think_end = None eos_token_ids: list[int] = [] @staticmethod def decode(_tokens: list[int]) -> str: return "hi" @staticmethod def encode(_text: str, add_special_tokens: bool = True) -> list[int]: return [0] class MockGroup: def rank(self) -> int: return 0 def size(self) -> int: return 1 def _run(tasks: Iterable[Task], send_after_ready: list[Task] | None = None): bound_instance = get_bound_mlx_ring_instance( instance_id=INSTANCE_1_ID, model_id=MODEL_A_ID, runner_id=RUNNER_1_ID, node_id=NODE_A, ) task_sender, task_receiver = mp_channel[Task]() _cancel_sender, cancel_receiver = mp_channel[TaskId]() on_event: Callable[[Event], None] | None = None if send_after_ready: _saw_running = False def _on_event(event: Event) -> None: nonlocal _saw_running if isinstance(event, RunnerStatusUpdated): if isinstance(event.runner_status, RunnerRunning): _saw_running = True elif _saw_running and isinstance(event.runner_status, RunnerReady): for t in send_after_ready: task_sender.send(t) on_event = _on_event event_sender = EventCollector(on_event=on_event) with task_sender: for t in tasks: task_sender.send(t) # worst monkeypatch known to man # this is some c++ nonsense task_receiver.close = nothin task_receiver.join = nothin with unittest.mock.patch( "exo.worker.runner.llm_inference.runner.mx.distributed.all_gather", make_nothin(mx.array([1])), ): runner = mlx_runner.Runner( bound_instance, event_sender, # pyright: ignore[reportArgumentType] task_receiver, cancel_receiver, ) runner.main() return event_sender.events def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch): events = _run( [INIT_TASK, LOAD_TASK, WARMUP_TASK, CHAT_TASK], send_after_ready=[SHUTDOWN_TASK], ) expected_chunk = ChunkGenerated( command_id=COMMAND_1_ID, chunk=TokenChunk( model=MODEL_A_ID, text="hi", token_id=0, finish_reason="stop", usage=None, stats=None, ), ) assert_events_equal( events, [ RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerIdle()), TaskStatusUpdated( task_id=INITIALIZATION_TASK_ID, task_status=TaskStatus.Running ), RunnerStatusUpdated( runner_id=RUNNER_1_ID, runner_status=RunnerConnecting() ), TaskAcknowledged(task_id=INITIALIZATION_TASK_ID), TaskStatusUpdated( task_id=INITIALIZATION_TASK_ID, task_status=TaskStatus.Complete ), RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerConnected()), TaskStatusUpdated(task_id=LOAD_TASK_ID, task_status=TaskStatus.Running), RunnerStatusUpdated( runner_id=RUNNER_1_ID, runner_status=RunnerLoading(layers_loaded=0, total_layers=32), ), TaskAcknowledged(task_id=LOAD_TASK_ID), TaskStatusUpdated(task_id=LOAD_TASK_ID, task_status=TaskStatus.Complete), RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerLoaded()), TaskStatusUpdated(task_id=WARMUP_TASK_ID, task_status=TaskStatus.Running), RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerWarmingUp()), TaskAcknowledged(task_id=WARMUP_TASK_ID), TaskStatusUpdated(task_id=WARMUP_TASK_ID, task_status=TaskStatus.Complete), RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerReady()), TaskStatusUpdated( task_id=CHAT_COMPLETION_TASK_ID, task_status=TaskStatus.Running ), RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerRunning()), TaskAcknowledged(task_id=CHAT_COMPLETION_TASK_ID), expected_chunk, TaskStatusUpdated( task_id=CHAT_COMPLETION_TASK_ID, task_status=TaskStatus.Complete ), # CHAT COMPLETION TASK SHOULD COMPLETE BEFORE RUNNER READY RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerReady()), TaskStatusUpdated(task_id=SHUTDOWN_TASK_ID, task_status=TaskStatus.Running), RunnerStatusUpdated( runner_id=RUNNER_1_ID, runner_status=RunnerShuttingDown() ), TaskAcknowledged(task_id=SHUTDOWN_TASK_ID), TaskStatusUpdated( task_id=SHUTDOWN_TASK_ID, task_status=TaskStatus.Complete ), # SPECIAL EXCEPTION FOR RUNNER SHUTDOWN RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerShutdown()), ], ) ================================================ FILE: src/exo/worker/tests/unittests/test_runner/test_glm_tool_parsing.py ================================================ """Tests for GLM tool call argument parsing regex.""" import regex as re # Replicate the regex patterns from runner.py to test them in isolation _func_name_regex = re.compile(r"^(.*?)", re.DOTALL) _func_arg_regex = re.compile( r"(.*?)(?:\n|\s)*(.*?)(?:|(?=)|$)", re.DOTALL, ) def _parse_args(text: str) -> list[tuple[str, str]]: """Extract (key, value) pairs from GLM tool call text.""" pairs = _func_arg_regex.findall(text) return [(k.strip(), v.strip()) for k, v in pairs] # pyright: ignore[reportAny] def _parse_func_name(text: str) -> str: """Extract function name from GLM tool call text.""" match = _func_name_regex.search(text) if match is None: raise ValueError(f"Could not parse function name: {text!r}") return match.group(1).strip() class TestGlmToolParsingWithClosingTags: """Tests for normal format with closing tags present.""" def test_single_argument(self): text = ( "get_weatherlocationSan Francisco" ) assert _parse_func_name(text) == "get_weather" pairs = _parse_args(text) assert pairs == [("location", "San Francisco")] def test_multiple_arguments(self): text = ( "searchquerypython" "limit10" ) assert _parse_func_name(text) == "search" pairs = _parse_args(text) assert pairs == [("query", "python"), ("limit", "10")] def test_arguments_with_whitespace_between(self): text = ( "fna\n1\n" "b 2" ) pairs = _parse_args(text) assert pairs == [("a", "1"), ("b", "2")] class TestGlmToolParsingMissingClosingTags: """Tests for format where closing tags are missing.""" def test_single_argument_no_closing(self): text = "get_weatherlocationSan Francisco" assert _parse_func_name(text) == "get_weather" pairs = _parse_args(text) assert pairs == [("location", "San Francisco")] def test_multiple_arguments_no_closing(self): text = ( "searchquerypython" "limit10" ) assert _parse_func_name(text) == "search" pairs = _parse_args(text) assert pairs == [("query", "python"), ("limit", "10")] def test_mixed_closing_tags(self): """First arg has closing tag, second does not.""" text = ( "fna1" "b2" ) pairs = _parse_args(text) assert pairs == [("a", "1"), ("b", "2")] def test_value_with_trailing_whitespace(self): text = "fnxhello world \n" pairs = _parse_args(text) assert pairs == [("x", "hello world")] def test_value_with_newlines_no_closing(self): text = "fndataline1\nline2" pairs = _parse_args(text) assert pairs == [("data", "line1\nline2")] class TestGlmToolParsingEdgeCases: """Edge case tests for GLM tool call parsing.""" def test_empty_value_with_closing(self): text = "fnempty" pairs = _parse_args(text) assert pairs == [("empty", "")] def test_value_with_json_content(self): text = 'fndata{"key": "value"}' pairs = _parse_args(text) assert pairs == [("data", '{"key": "value"}')] def test_value_with_json_no_closing(self): text = 'fndata{"key": "value"}' pairs = _parse_args(text) assert pairs == [("data", '{"key": "value"}')] ================================================ FILE: src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py ================================================ from collections.abc import Generator from exo.api.types import FinishReason from exo.shared.types.worker.runner_response import ( GenerationResponse, ToolCallResponse, ) from exo.worker.runner.llm_inference.model_output_parsers import parse_gpt_oss # Token IDs from mlx-community/gpt-oss-20b-MXFP4-Q8 tokenizer. # These are stable since they come from the model's vocabulary. _CHANNEL = 200005 # <|channel|> _START = 200006 # <|start|> _MESSAGE = 200008 # <|message|> _CALL = 200012 # <|call|> _END = 200007 # <|end|> _ASSISTANT = 173781 # "assistant" # fmt: off # " to=functions.get_current_weather<|channel|>commentary json<|message|>{\"location\": \"Tokyo\"}<|call|>" FORMAT_A_TOKENS: list[tuple[int, str]] = [ (316, " to"), (28, "="), (44580, "functions"), (775, ".get"), (23981, "_current"), (170154, "_weather"), (_CHANNEL, "<|channel|>"), (12606, "comment"), (815, "ary"), (5701, " json"), (_MESSAGE, "<|message|>"), (10848, '{"'), (7693, "location"), (1243, '":'), (392, ' "'), (173844, "Tokyo"), (18583, '"}'), (_CALL, "<|call|>"), ] # "<|channel|>commentary to=functions.get_current_weather json<|message|>{\"location\": \"Tokyo\"}<|call|>" FORMAT_B_TOKENS: list[tuple[int, str]] = [ (_CHANNEL, "<|channel|>"), (12606, "comment"), (815, "ary"), (316, " to"), (28, "="), (44580, "functions"), (775, ".get"), (23981, "_current"), (170154, "_weather"), (5701, " json"), (_MESSAGE, "<|message|>"), (10848, '{"'), (7693, "location"), (1243, '":'), (392, ' "'), (173844, "Tokyo"), (18583, '"}'), (_CALL, "<|call|>"), ] # "<|channel|>analysis<|message|>Let me think...<|end|><|start|>assistant<|channel|>commentary to=functions.X ..." # Full analysis-then-tool-call as the model actually generates it. THINKING_THEN_TOOL_TOKENS: list[tuple[int, str]] = [ (_CHANNEL, "<|channel|>"), (35644, "analysis"), (_MESSAGE, "<|message|>"), (12845, "Let"), (668, " me"), (2411, " think"), (1078, " about"), (495, " this"), (13, "."), (_END, "<|end|>"), # Model generates a new message header for the tool call: (_START, "<|start|>"), (_ASSISTANT, "assistant"), *FORMAT_B_TOKENS, ] # fmt: on def _make_gen_responses( tokens: list[tuple[int, str]], last_finish_reason: FinishReason = "stop", ) -> list[GenerationResponse]: """Build GenerationResponse list from (token_id, text) pairs.""" responses: list[GenerationResponse] = [] for i, (tid, text) in enumerate(tokens): is_last = i == len(tokens) - 1 responses.append( GenerationResponse( text=text, token=tid, finish_reason=last_finish_reason if is_last else None, usage=None, ) ) return responses def _collect( tokens: list[tuple[int, str]], last_finish_reason: FinishReason = "stop", ) -> list[GenerationResponse | ToolCallResponse]: """Feed tokens through parse_gpt_oss and collect all yielded responses.""" def _gen() -> Generator[GenerationResponse, None, None]: yield from _make_gen_responses(tokens, last_finish_reason) return list(x for x in parse_gpt_oss(_gen()) if x is not None) def _get_tool_call( results: list[GenerationResponse | ToolCallResponse], ) -> ToolCallResponse: """Extract the single ToolCallResponse from results.""" tool_calls = [r for r in results if isinstance(r, ToolCallResponse)] assert len(tool_calls) == 1, f"Expected 1 ToolCallResponse, got {len(tool_calls)}" return tool_calls[0] class TestParseGptOssRecipientPlacement: """Both Harmony recipient placements must produce identical tool calls.""" def test_format_a_yields_tool_call(self): results = _collect(FORMAT_A_TOKENS) tc = _get_tool_call(results) assert tc.tool_calls[0].name == "get_current_weather" assert '"location"' in tc.tool_calls[0].arguments assert "Tokyo" in tc.tool_calls[0].arguments def test_format_b_yields_tool_call(self): results = _collect(FORMAT_B_TOKENS) tc = _get_tool_call(results) assert tc.tool_calls[0].name == "get_current_weather" assert '"location"' in tc.tool_calls[0].arguments assert "Tokyo" in tc.tool_calls[0].arguments def test_both_formats_produce_identical_tool_calls(self): tc_a = _get_tool_call(_collect(FORMAT_A_TOKENS)) tc_b = _get_tool_call(_collect(FORMAT_B_TOKENS)) assert tc_a.tool_calls[0].name == tc_b.tool_calls[0].name assert tc_a.tool_calls[0].arguments == tc_b.tool_calls[0].arguments class TestParseGptOssThinkingThenToolCall: """Analysis (thinking) followed by a tool call must yield both.""" def test_thinking_then_tool_call(self): results = _collect(THINKING_THEN_TOOL_TOKENS) # Thinking tokens should have is_thinking=True and no tags thinking_responses = [ r for r in results if isinstance(r, GenerationResponse) and r.is_thinking ] thinking_text = "".join(r.text for r in thinking_responses) assert "Let me think about this." in thinking_text assert "" not in thinking_text assert "" not in thinking_text # Non-thinking tokens should have is_thinking=False non_thinking = [ r for r in results if isinstance(r, GenerationResponse) and not r.is_thinking ] non_thinking_text = "".join(r.text for r in non_thinking) assert "" not in non_thinking_text # And the tool call tc = _get_tool_call(results) assert tc.tool_calls[0].name == "get_current_weather" assert "Tokyo" in tc.tool_calls[0].arguments # fmt: off # Truncated tool call: recipient + channel + message + partial args, no <|call|> TRUNCATED_TOOL_CALL_TOKENS: list[tuple[int, str]] = [ (316, " to"), (28, "="), (44580, "functions"), (775, ".get"), (23981, "_current"), (170154, "_weather"), (_CHANNEL, "<|channel|>"), (12606, "comment"), (815, "ary"), (5701, " json"), (_MESSAGE, "<|message|>"), (10848, '{"'), (7693, "location"), (1243, '":'), (392, ' "'), (173844, "Tokyo"), # No <|call|> — generation truncated here ] # Plain text tokens (no tool call) PLAIN_TEXT_TOKENS: list[tuple[int, str]] = [ (_CHANNEL, "<|channel|>"), (35644, "analysis"), (_MESSAGE, "<|message|>"), (12845, "Let"), (668, " me"), (2411, " think"), (1078, " about"), (495, " this"), (13, "."), (_END, "<|end|>"), (_START, "<|start|>"), (_ASSISTANT, "assistant"), (_CHANNEL, "<|channel|>"), (12606, "comment"), (815, "ary"), (_MESSAGE, "<|message|>"), (9906, "Hello"), (14, ","), (2989, " world"), ] # fmt: on class TestParseGptOssMaxTokensTruncation: """Truncated tool calls must still yield finish_reason.""" def test_truncated_tool_call_yields_finish_reason(self): results = _collect(TRUNCATED_TOOL_CALL_TOKENS, last_finish_reason="length") gen_responses = [r for r in results if isinstance(r, GenerationResponse)] finish_reasons = [ r.finish_reason for r in gen_responses if r.finish_reason is not None ] assert "length" in finish_reasons def test_truncated_tool_call_emits_partial_args(self): results = _collect(TRUNCATED_TOOL_CALL_TOKENS, last_finish_reason="length") gen_responses = [r for r in results if isinstance(r, GenerationResponse)] last = [r for r in gen_responses if r.finish_reason is not None][-1] assert len(last.text) > 0 def test_truncated_plain_text_still_works(self): results = _collect(PLAIN_TEXT_TOKENS, last_finish_reason="length") gen_responses = [r for r in results if isinstance(r, GenerationResponse)] finish_reasons = [ r.finish_reason for r in gen_responses if r.finish_reason is not None ] assert "length" in finish_reasons # Verify non-empty text was yielded (delta text differs from raw token text # due to Harmony encoding, so we just check something was emitted) all_text = "".join(r.text for r in gen_responses) assert len(all_text) > 0 ================================================ FILE: src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py ================================================ """Tests for parse_tool_calls generator, especially unclosed tool call handling.""" import json from collections.abc import Generator from typing import Any from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse from exo.worker.runner.llm_inference.model_output_parsers import parse_tool_calls from exo.worker.runner.llm_inference.tool_parsers import make_mlx_parser def _make_responses( texts: list[str], finish_on_last: bool = True, ) -> Generator[GenerationResponse]: """Create a sequence of GenerationResponses from text strings.""" for i, text in enumerate(texts): is_last = i == len(texts) - 1 yield GenerationResponse( text=text, token=i, finish_reason="stop" if (is_last and finish_on_last) else None, usage=None, ) def _dummier_parser(text: str) -> dict[str, Any]: return {"name": "test_fn", "arguments": {"arg": text}} _dummy_parser = make_mlx_parser("", "", _dummier_parser) class TestParseToolCalls: """Tests for parse_tool_calls generator.""" def test_closed_tool_call_works_normally(self): """Normal tool call flow should not be affected.""" texts = ["", "test_fn", ""] results = list( parse_tool_calls( _make_responses(texts, finish_on_last=False), _dummy_parser, tools=None, ) ) assert len(results) == 1 assert isinstance(results[0], ToolCallResponse) def test_no_tool_call_passes_through(self): """Responses without tool calls should pass through unchanged.""" texts = ["Hello", " world"] results = list( parse_tool_calls( _make_responses(texts), _dummy_parser, tools=None, ) ) assert len(results) == 2 assert all(isinstance(r, GenerationResponse) for r in results) r0 = results[0] r1 = results[1] assert isinstance(r0, GenerationResponse) assert isinstance(r1, GenerationResponse) assert r0.text == "Hello" assert r1.text == " world" assert r1.finish_reason == "stop" def test_failed_parse_yields_text(self): """When tool call parsing fails, the text should be yielded as-is.""" def _failing_parser(text: str) -> dict[str, Any]: raise ValueError("parse failed") texts = ["", "bad content", ""] results = list( parse_tool_calls( _make_responses(texts, finish_on_last=False), make_mlx_parser("", "", _failing_parser), tools=None, ) ) assert len(results) == 1 assert isinstance(results[0], GenerationResponse) assert results[0].text == "bad content" def test_tool_schema_coerces_string_arguments_to_expected_types(self): """Tool argument values should be coerced using provided JSON schema.""" def _parser_with_string_args(_text: str) -> dict[str, Any]: return { "name": "process", "arguments": { "action": "output", "id": "0", "verbose": "true", "temperature": "0.75", }, } tools = [ { "type": "function", "function": { "name": "process", "description": "Manage background processes", "parameters": { "type": "object", "properties": { "action": {"type": "string"}, "id": {"type": "integer"}, "verbose": {"type": "boolean"}, "temperature": {"type": "number"}, }, "required": ["action"], }, }, } ] results = list( parse_tool_calls( _make_responses(["", "process", ""]), make_mlx_parser( "", "", _parser_with_string_args ), tools, ) ) assert len(results) == 1 assert isinstance(results[0], ToolCallResponse) args = json.loads(results[0].tool_calls[0].arguments) # pyright: ignore[reportAny] assert args == { "action": "output", "id": 0, "verbose": True, "temperature": 0.75, } def test_schema_coercion_skips_unknown_tools(self): """If no matching tool schema exists, arguments should remain unchanged.""" def _parser_with_string_id(_text: str) -> dict[str, Any]: return { "name": "process", "arguments": {"action": "output", "id": "0"}, } tools = [ { "type": "function", "function": { "name": "different_tool", "parameters": { "type": "object", "properties": {"id": {"type": "integer"}}, }, }, } ] results = list( parse_tool_calls( _make_responses(["", "process", ""]), make_mlx_parser("", "", _parser_with_string_id), tools, ) ) assert len(results) == 1 assert isinstance(results[0], ToolCallResponse) args = json.loads(results[0].tool_calls[0].arguments) # pyright: ignore[reportAny] assert args == {"action": "output", "id": "0"} ================================================ FILE: src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py ================================================ import multiprocessing as mp from typing import cast import anyio import pytest from exo.shared.models.model_cards import ModelId from exo.shared.types.chunks import ErrorChunk from exo.shared.types.common import CommandId, NodeId from exo.shared.types.events import ChunkGenerated, Event, RunnerStatusUpdated from exo.shared.types.tasks import Task, TaskId, TextGeneration from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams from exo.shared.types.worker.instances import BoundInstance, InstanceId from exo.shared.types.worker.runners import RunnerFailed, RunnerId from exo.utils.channels import channel, mp_channel from exo.worker.runner.runner_supervisor import RunnerSupervisor from exo.worker.tests.unittests.conftest import get_bound_mlx_ring_instance class _DeadProcess: exitcode = -6 def start(self) -> None: return None def is_alive(self) -> bool: return False def join(self, _timeout: float | None = None) -> None: return None def terminate(self) -> None: return None def kill(self) -> None: return None @pytest.mark.asyncio async def test_check_runner_emits_error_chunk_for_inflight_text_generation() -> None: event_sender, event_receiver = channel[Event]() task_sender, _ = mp_channel[Task]() cancel_sender, _ = mp_channel[TaskId]() _, ev_recv = mp_channel[Event]() bound_instance: BoundInstance = get_bound_mlx_ring_instance( instance_id=InstanceId("instance-a"), model_id=ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit"), runner_id=RunnerId("runner-a"), node_id=NodeId("node-a"), ) supervisor = RunnerSupervisor( shard_metadata=bound_instance.bound_shard, bound_instance=bound_instance, runner_process=cast("mp.Process", cast(object, _DeadProcess())), initialize_timeout=400, _ev_recv=ev_recv, _task_sender=task_sender, _event_sender=event_sender, _cancel_sender=cancel_sender, ) command_id = CommandId("cmd-a") task = TextGeneration( task_id=TaskId("task-a"), instance_id=bound_instance.instance.instance_id, command_id=command_id, task_params=TextGenerationTaskParams( model=bound_instance.bound_shard.model_card.model_id, input=[InputMessage(role="user", content="hi")], stream=True, ), ) supervisor.in_progress[task.task_id] = task supervisor.shutdown = lambda: None await supervisor._check_runner(RuntimeError("boom")) # pyright: ignore[reportPrivateUsage] got_chunk = await event_receiver.receive() got_status = await event_receiver.receive() assert isinstance(got_chunk, ChunkGenerated) assert got_chunk.command_id == command_id assert isinstance(got_chunk.chunk, ErrorChunk) assert "Runner shutdown before completing command" in got_chunk.chunk.error_message assert isinstance(got_status, RunnerStatusUpdated) assert isinstance(got_status.runner_status, RunnerFailed) event_sender.close() with anyio.move_on_after(0.1): await event_receiver.aclose() ================================================ FILE: tests/auto_bench.sh ================================================ #!/usr/bin/env bash [ $# -lt 1 ] && { echo "Usage: $0 host1 [host2 ...]" exit 1 } [ -z "$(git status --porcelain)" ] || { echo "Uncommitted changes" exit 1 } commit=$(git rev-parse HEAD) git fetch -q origin git branch -r --contains "$commit" | grep -qE '^\s*origin/' || { echo "Not pushed to origin" exit 1 } hosts=("$@") for host; do ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \ "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" & done wait cleanup() { for host in "${hosts[@]}"; do ssh -T -o BatchMode=yes "$host@$host" "pkill -f bin/exo" & done sleep 1 jobs -pr | xargs -r kill 2>/dev/null || true } trap 'cleanup' EXIT INT TERM for host; do ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \ "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" &>/dev/null & done for host; do echo "Waiting for $host..." 1>&2 until curl -sf "http://$host:52415/models" &>/dev/null; do sleep 1; done done echo "Waiting 30s for cluster setup" 1>&2 sleep 30 echo "EXO loaded" 1>&2 bench_runner="${hosts[0]}" mkdir -p "./bench/$commit" nix run .#exo-get-all-models-on-cluster -- "$bench_runner" | while IFS= read -r model; do echo "running bench for $model" 1>&2 ssh -Tn -o BatchMode=yes -o ServerAliveInterval=30 "$bench_runner@$bench_runner" "/nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit#exo-bench -- --model $model --pp 128 4096 --tg 128 --stdout --skip-tensor-ring" >>"./bench/$commit/${model//\//--}.json" echo done ================================================ FILE: tests/eval_tool_calls.sh ================================================ #!/usr/bin/env bash [ $# -lt 1 ] && { echo "Usage: $0 host1 [host2 ...]" exit 1 } [ -z "$(git status --porcelain)" ] || { echo "Uncommitted changes" exit 1 } commit=$(git rev-parse HEAD) git fetch -q origin git branch -r --contains "$commit" | grep -qE '^\s*origin/' || { echo "Not pushed to origin" exit 1 } hosts=("$@") cleanup() { for host in "${hosts[@]}"; do ssh -T -o BatchMode=yes "$host@$host" "pkill -f bin/exo" & done sleep 1 jobs -pr | xargs -r kill 2>/dev/null || true } trap 'cleanup' EXIT INT TERM for host; do ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \ "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" & done wait for host; do ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \ "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" &>/dev/null & done for host; do echo "Waiting for $host..." 1>&2 until curl -sf "http://$host:52415/models" &>/dev/null; do sleep 1; done done echo "Waiting 30s for cluster setup" 1>&2 sleep 30 echo "EXO loaded" 1>&2 eval_runner="${hosts[0]}" mkdir -p "./bench/$commit" nix run .#exo-get-all-models-on-cluster -- "$eval_runner" | while IFS= read -r model; do echo "running eval for $model" 1>&2 ssh -Tn -o BatchMode=yes -o ServerAliveInterval=30 "$eval_runner@$eval_runner" \ "/nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit#exo-eval-tool-calls -- --model $model --stdout" \ >>"./bench/$commit/${model//\//--}-eval.json" echo done ================================================ FILE: tests/get_all_models_on_cluster.py ================================================ #!/usr/bin/env python3 # pyright: reportAny=false import json import subprocess import sys from typing import Any, cast from urllib.request import urlopen h = sys.argv[1] if len(sys.argv) > 1 else sys.exit(f"USAGE: {sys.argv[0]} host") ts = subprocess.run( ["tailscale", "status"], check=True, text=True, capture_output=True ).stdout.splitlines() ip = next( (sl[0] for line in ts if len(sl := line.split()) >= 2 if sl[1] == h), None ) or sys.exit(f"{h} not found in tailscale") with urlopen(f"http://{ip}:52415/state", timeout=5) as r: data = json.loads(r.read()).get("downloads", {}) def mid(x: dict[str, Any]) -> str | None: for k in ( "DownloadCompleted", "shardMetadata", "PipelineShardMetadata", "modelCard", "modelId", ): x = x.get(k, {}) return cast(str | None, x if x != {} else None) common = set[str].intersection( *[{m for d in nid if (m := mid(d))} for nid in data.values()] ) for c in common: print(c) ================================================ FILE: tests/headless_runner.py ================================================ import socket from typing import Literal import anyio from fastapi import FastAPI from fastapi.responses import Response, StreamingResponse from hypercorn import Config from hypercorn.asyncio import serve # pyright: ignore[reportUnknownVariableType] from loguru import logger from pydantic import BaseModel from exo.shared.constants import EXO_MODELS_DIR from exo.shared.models.model_cards import ModelCard, ModelId from exo.shared.types.chunks import TokenChunk from exo.shared.types.commands import CommandId from exo.shared.types.common import Host, NodeId from exo.shared.types.events import ChunkGenerated, Event, RunnerStatusUpdated from exo.shared.types.tasks import ( ConnectToGroup, LoadModel, Shutdown, StartWarmup, Task, TextGeneration, ) from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams from exo.shared.types.worker.instances import ( BoundInstance, Instance, InstanceId, MlxJacclInstance, MlxRingInstance, ) from exo.shared.types.worker.runners import ( RunnerFailed, RunnerId, RunnerShutdown, ShardAssignments, ) from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata from exo.utils.channels import channel, mp_channel from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer from exo.worker.runner.bootstrap import entrypoint class Tests(BaseModel): # list[hostname, ip addr] devs: list[list[str]] ibv_devs: list[list[str | None]] | None model_id: ModelId kind: Literal["ring", "jaccl", "both"] iid = InstanceId("im testing here") async def main(): logger.info("starting cool server majig") cfg = Config() cfg.bind = "0.0.0.0:52414" # nb: shared.logging needs updating if any of this changes cfg.accesslog = "-" cfg.errorlog = "-" ev = anyio.Event() app = FastAPI() app.post("/run_test")(run_test) app.post("/kill")(lambda: kill(ev)) app.get("/tb_detection")(tb_detection) app.get("/models")(list_models) await serve( app, # type: ignore cfg, shutdown_trigger=lambda: ev.wait(), ) def kill(ev: anyio.Event): ev.set() return Response(status_code=204) async def tb_detection(): send, recv = channel[GatheredInfo]() ig = InfoGatherer(send) with anyio.move_on_after(1): await ig._monitor_system_profiler_thunderbolt_data() # pyright: ignore[reportPrivateUsage] with recv: return recv.collect() def list_models(): sent = set[str]() for path in EXO_MODELS_DIR.rglob("model-*.safetensors"): if "--" not in path.parent.name: continue name = path.parent.name.replace("--", "/") if name in sent: continue sent.add(name) yield ModelId(path.parent.name.replace("--", "/")) async def run_test(test: Tests): weird_hn = socket.gethostname() for dev in test.devs: if weird_hn.startswith(dev[0]) or dev[0].startswith(weird_hn): hn = dev[0] break else: raise ValueError(f"{weird_hn} not in {test.devs}") async def run(): logger.info(f"testing {test.model_id}") instances: list[Instance] = [] if test.kind in ["ring", "both"]: i = await ring_instance(test, hn) if i is None: yield "no model found" return instances.append(i) if test.kind in ["jaccl", "both"]: i = await jaccl_instance(test) if i is None: yield "no model found" return instances.append(i) for instance in instances: recv = await execute_test(test, instance, hn) str_out = "" for item in recv: if isinstance(item, ChunkGenerated): assert isinstance(item.chunk, TokenChunk) str_out += item.chunk.text if isinstance(item, RunnerStatusUpdated) and isinstance( item.runner_status, (RunnerFailed, RunnerShutdown) ): yield str_out + "\n" yield item.model_dump_json() + "\n" return StreamingResponse(run()) async def ring_instance(test: Tests, hn: str) -> Instance | None: hbn = [Host(ip="198.51.100.0", port=52417) for _ in test.devs] world_size = len(test.devs) for i in range(world_size): if test.devs[i][0] == hn: hn = test.devs[i][0] hbn[(i - 1) % world_size] = Host(ip=test.devs[i - 1][1], port=52417) hbn[(i + 1) % world_size] = Host(ip=test.devs[i + 1][1], port=52417) hbn[i] = Host(ip="0.0.0.0", port=52417) break else: raise ValueError(f"{hn} not in {test.devs}") card = await ModelCard.load(test.model_id) instance = MlxRingInstance( instance_id=iid, ephemeral_port=52417, hosts_by_node={NodeId(hn): hbn}, shard_assignments=ShardAssignments( model_id=test.model_id, node_to_runner={NodeId(host[0]): RunnerId(host[0]) for host in test.devs}, runner_to_shard={ RunnerId(test.devs[i][0]): PipelineShardMetadata( model_card=card, device_rank=i, world_size=world_size, start_layer=(card.n_layers // world_size) * i, end_layer=min( card.n_layers, (card.n_layers // world_size) * (i + 1) ), n_layers=min(card.n_layers, (card.n_layers // world_size) * (i + 1)) - (card.n_layers // world_size) * i, ) for i in range(world_size) }, ), ) return instance async def execute_test(test: Tests, instance: Instance, hn: str) -> list[Event]: world_size = len(test.devs) commands: list[Task] = [ (LoadModel(instance_id=iid)), (StartWarmup(instance_id=iid)), ( TextGeneration( task_params=TextGenerationTaskParams( model=test.model_id, instructions="You are a helpful assistant", input=[ InputMessage( role="user", content="What is the capital of France?" ) ], ), command_id=CommandId("yo"), instance_id=iid, ) ), (Shutdown(runner_id=RunnerId(hn), instance_id=iid)), ] if world_size > 1: commands.insert(0, ConnectToGroup(instance_id=iid)) bound_instance = BoundInstance( instance=instance, bound_runner_id=RunnerId(hn), bound_node_id=NodeId(hn) ) ev_send, _ev_recv = mp_channel[Event]() task_send, task_recv = mp_channel[Task]() for command in commands: task_send.send(command) entrypoint( bound_instance, ev_send, task_recv, logger, ) # TODO(evan): return ev_recv.collect() return [] async def jaccl_instance(test: Tests) -> MlxJacclInstance | None: card = await ModelCard.load(test.model_id) world_size = len(test.devs) assert test.ibv_devs return MlxJacclInstance( instance_id=iid, jaccl_devices=test.ibv_devs, # rank 0 is always coordinator jaccl_coordinators={ NodeId(host[0]): test.devs[0][1] + ":52417" for host in test.devs }, shard_assignments=ShardAssignments( model_id=test.model_id, node_to_runner={NodeId(host[0]): RunnerId(host[0]) for host in test.devs}, runner_to_shard={ RunnerId(host[0]): TensorShardMetadata( model_card=card, device_rank=i, world_size=world_size, start_layer=0, end_layer=card.n_layers, n_layers=card.n_layers, ) for i, host in enumerate(test.devs) }, ), ) if __name__ == "__main__": anyio.run(main) ================================================ FILE: tests/run_exo_on.sh ================================================ #!/usr/bin/env bash set -euo pipefail [ $# -lt 1 ] && { echo "Usage: $0 host1 [host2 ...]" exit 1 } [ -z "$(git status --porcelain)" ] || { echo "Uncommitted changes" exit 1 } commit=$(git rev-parse HEAD) git fetch -q origin git branch -r --contains "$commit" | grep -qE '^\s*origin/' || { echo "Not pushed to origin" exit 1 } echo "Deploying $commit to $# hosts..." hosts=("$@") cleanup() { for host in "${hosts[@]}"; do ssh -T -o BatchMode=yes "$host@$host" "pkill -f bin/exo" & done wait jobs -pr | xargs -r kill 2>/dev/null || true } trap 'cleanup' EXIT INT TERM colours=($'\e[31m' $'\e[32m' $'\e[33m' $'\e[34m') reset=$'\e[0m' i=0 for host; do colour=${colours[i++ % 4]} ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \ "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" |& awk -v p="${colour}[${host}]${reset}" '{ print p $0; fflush() }' & done for host; do echo "Waiting for $host..." until curl -sf "http://$host:52415/models" &>/dev/null; do sleep 1; done done wait ================================================ FILE: tests/start_distributed_test.py ================================================ #!/usr/bin/env python3 import itertools import json import subprocess import sys from concurrent.futures import ThreadPoolExecutor from typing import Any, cast from urllib.request import Request, urlopen if not (args := sys.argv[1:]): sys.exit( f"USAGE: {sys.argv[0]} [host1] [host2] ...\nkind is optional, and should be jaccl or ring" ) kind = args[0] if args[0] in ("jaccl", "ring") else "both" hosts = args[1:] if kind != "both" else args ts = subprocess.run( ["tailscale", "status"], check=True, text=True, capture_output=True ).stdout.splitlines() ip = {sl[1]: sl[0] for line in ts if len(sl := line.split()) >= 2} ips = [ip[h] for h in hosts] devs = [[h, ip[h]] for h in hosts] n = len(hosts) def get_tb(a: str) -> list[dict[str, Any]]: with urlopen(f"http://{a}:52414/tb_detection", timeout=5) as r: # pyright: ignore[reportAny] return json.loads(r.read()) # pyright: ignore[reportAny] def get_models(a: str) -> set[str]: with urlopen(f"http://{a}:52414/models", timeout=5) as r: # pyright: ignore[reportAny] return set(json.loads(r.read())) # pyright: ignore[reportAny] def run(h: str, a: str, body: bytes) -> None: with urlopen( Request( f"http://{a}:52414/run_test", data=body, method="POST", headers={"Content-Type": "application/json"}, ), timeout=300, ) as r: # pyright: ignore[reportAny] for line in r.read().decode(errors="replace").splitlines(): # pyright: ignore[reportAny] print(f"\n{h}@{a}: {line}", flush=True) with ThreadPoolExecutor(n) as exctr: if kind in ("jaccl", "both"): payloads = list(exctr.map(get_tb, ips)) u2e = { ident["domainUuid"]: (i, ident["rdmaInterface"]) for i, p in enumerate(payloads) for d in p for ident in cast( list[dict[str, str]], d.get("MacThunderboltIdentifiers", {}).get("idents", []), # pyright: ignore[reportAny] ) } edges = { (u2e[s][0], u2e[t][0]): u2e[t][1] for p in payloads for d in p for c in d.get("MacThunderboltConnections", {}).get("conns", []) # pyright: ignore[reportAny] if (s := c["sourceUuid"]) in u2e and (t := c["sinkUuid"]) in u2e # pyright: ignore[reportAny] } ibv_devs = [[edges.get((i, j)) for j in range(n)] for i in range(n)] else: ibv_devs = None models = set[str].intersection(*exctr.map(get_models, ips)) print("\n") print("=" * 70) print(f"Starting test with {models}") print("=" * 70) print("\n") for model in models: body = json.dumps( {"devs": devs, "model_id": model, "ibv_devs": ibv_devs, "kind": kind} ).encode() list(exctr.map(run, hosts, ips, itertools.repeat(body))) ================================================ FILE: tmp/config_examples/claude_code.sh ================================================ #!/bin/bash # Run Claude Code against a local exo cluster! (Here, GPT OSS 120B) ANTHROPIC_BASE_URL="http://localhost:52415/" \ ANTHROPIC_AUTH_TOKEN="dummy" \ ANTHROPIC_MODEL="mlx-community/gpt-oss-120b-MXFP4-Q8" \ ANTHROPIC_SMALL_FAST_MODEL="mlx-community/gpt-oss-120b-MXFP4-Q8" \ CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1 \ claude ================================================ FILE: tmp/config_examples/opencode.json ================================================ { "$schema": "https://opencode.ai/config.json", "model": "exo/mlx-community/gpt-oss-120b-MXFP4-Q8", "provider": { "exo": { "api": "http://localhost:52415/v1", "models": { "mlx-community/gpt-oss-120b-MXFP4-Q8": { "name": "GPT OSS 120B", "limit": { "context": 32768, "output": 8192 } } } } } } ================================================ FILE: tmp/gen_card.py ================================================ """ Generates inference model cards for EXO. Usage: uv run tmp/gen_card.py mlx-community/my_cool_model-8bit [repo-id/model-id-2] [...] Model Cards require cleanup for family & quantization data """ import sys import anyio from exo.shared.models.model_cards import ModelCard, ModelId async def main(): if len(sys.argv) == 1: print(f"USAGE: {sys.argv[0]} repo-id/model-id-1 [repo-id/model-id-2] [...]") quit(1) print("Remember! Model Cards require cleanup for family & quantization data") for arg in sys.argv[1:]: mid = ModelId(arg) mc = await ModelCard.fetch_from_hf(mid) await mc.save( anyio.Path(__file__).parent.parent / "resources" / "inference_model_cards" / (mid.normalize() + ".toml") ) if __name__ == "__main__": anyio.run(main) ================================================ FILE: tmp/prompt.txt ================================================ Summarise this Wikipedia article for me: Transition from Republic to Empire Augustus of Prima Porta Rome had begun expanding shortly after the founding of the Roman Republic in the 6th century BC, though not outside the Italian Peninsula until the 3rd century BC. The Republic was not a nation-state in the modern sense, but a network of self-ruled towns (with varying degrees of independence from the Senate) and provinces administered by military commanders. It was governed by annually elected magistrates (Roman consuls above all) in conjunction with the Senate.[22] The 1st century BC was a time of political and military upheaval, which ultimately led to rule by emperors.[23][24][25] The consuls' military power rested in the Roman legal concept of imperium, meaning "command" (typically in a military sense).[26] Occasionally, successful consuls or generals were given the honorary title imperator (commander); this is the origin of the word emperor, since this title was always bestowed to the early emperors.[27][g] Rome suffered a long series of internal conflicts, conspiracies, and civil wars from the late second century BC (see Crisis of the Roman Republic) while greatly extending its power beyond Italy. In 44 BC Julius Caesar was briefly perpetual dictator before being assassinated by a faction that opposed his concentration of power. This faction was driven from Rome and defeated at the Battle of Philippi in 42 BC by Mark Antony and Caesar's adopted son Octavian. Antony and Octavian divided the Roman world between them, but this did not last long. Octavian's forces defeated those of Mark Antony and Cleopatra at the Battle of Actium in 31 BC. In 27 BC the Senate gave him the title Augustus ("venerated") and made him princeps ("foremost") with proconsular imperium, thus beginning the Principate, the first epoch of Roman imperial history. Although the republic stood in name, Augustus had all meaningful authority.[29] During his 40-year rule, a new constitutional order emerged so that, upon his death, Tiberius would succeed him as the new de facto monarch.[30] Pax Romana Main article: Pax Romana The so-called "Five Good Emperors" of 96–180 AD Nerva (r. 96–98) Trajan (r. 98–117) Hadrian (r. 117–138) Antoninus Pius (r. 138–161) Marcus Aurelius (r. 161–180) The 200 years that began with Augustus's rule are traditionally regarded as the Pax Romana ("Roman Peace"). The cohesion of the empire was furthered by a degree of social stability and economic prosperity that Rome had never before experienced. Uprisings in the provinces were infrequent and put down "mercilessly and swiftly".[31] The success of Augustus in establishing principles of dynastic succession was limited by his outliving a number of talented potential heirs. The Julio-Claudian dynasty lasted for four more emperors—Tiberius, Caligula, Claudius, and Nero—before it yielded in 69 AD to the strife-torn Year of the Four Emperors, from which Vespasian emerged as the victor. Vespasian became the founder of the brief Flavian dynasty, followed by the Nerva–Antonine dynasty which produced the "Five Good Emperors": Nerva, Trajan, Hadrian, Antoninus Pius, and Marcus Aurelius.[32] Among the so-called “Five Good Emperors,” Hadrian (r. 117–138) is particularly noted for consolidating the empire’s frontiers and embarking on ambitious building projects throughout the provinces.[33] In Judaea, which had long been the center of Jewish national and religious life, his reign marked a decisive turning point. After earlier Jewish resistance to Roman rule, Hadrian visited the region in 129/130 CE and refounded Jerusalem as the Roman colony Aelia Capitolina, naming it after his family (Aelius) and the Capitoline Triad.[34] The refoundation overlaid the destroyed Jewish city with a new Roman urban plan, and included the construction of a Temple to Jupiter on the site of the former Jewish Temple.[35] Later tradition and archaeological evidence also indicate a Temple of Venus near the site of the Holy Sepulchre.[36] Hadrian’s measures, combined with restrictions on Jewish practices, helped spark the Bar Kokhba Revolt (132–135 CE). After crushing the uprising, Roman forces expelled most Jews from Jerusalem, barring their entry except on certain days, and rebuilt the city as a statement of imperial power and domination.[33] Most scholars consider Hadrianic Aelia to have been unwalled, with free-standing gate complexes (such as the northern gate beneath today’s Damascus Gate) rather than a continuous defensive circuit.[37] Transition from classical to late antiquity Main articles: Later Roman Empire and Fall of the Western Roman Empire See also: Barbarian kingdoms and Byzantine Empire The Barbarian invasions consisted of the movement of (mainly) ancient Germanic peoples into Roman territory. Historically, this event marked the transition between classical antiquity and the Middle Ages. In the view of contemporary Greek historian Cassius Dio, the accession of Commodus in 180 marked the descent "from a kingdom of gold to one of rust and iron",[38] a comment which has led some historians, notably Edward Gibbon, to take Commodus' reign as the beginning of the Empire's decline.[39][40] In 212, during the reign of Caracalla, Roman citizenship was granted to all freeborn inhabitants of the empire. The Severan dynasty was tumultuous; an emperor's reign was ended routinely by his murder or execution and, following its collapse, the Empire was engulfed by the Crisis of the Third Century, a period of invasions, civil strife, economic disorder, and plague.[41] In defining historical epochs, this crisis sometimes marks the transition from Classical to Late Antiquity. Aurelian (r. 270–275) stabilised the empire militarily and Diocletian reorganised and restored much of it in 285.[42] Diocletian's reign brought the empire's most concerted effort against the perceived threat of Christianity, the "Great Persecution".[43] Diocletian divided the empire into four regions, each ruled by a separate tetrarch.[44] Confident that he fixed the disorder plaguing Rome, he abdicated along with his co-emperor, but the Tetrarchy collapsed shortly after. Order was eventually restored by Constantine the Great, who became the first emperor to convert to Christianity, and who established Constantinople as the new capital of the Eastern Empire. During the decades of the Constantinian and Valentinian dynasties, the empire was divided along an east–west axis, with dual power centres in Constantinople and Rome. Julian, who under the influence of his adviser Mardonius attempted to restore Classical Roman and Hellenistic religion, only briefly interrupted the succession of Christian emperors. Theodosius I, the last emperor to rule over both East and West, died in 395 after making Christianity the state religion.[45] The Roman Empire by 476, noting western and eastern divisions The administrative divisions of the Roman Empire in 395 AD Fall in the West and survival in the East The Western Roman Empire began to disintegrate in the early 5th century. The Romans fought off all invaders, most famously Attila,[46] but the empire had assimilated so many Germanic peoples of dubious loyalty to Rome that the empire started to dismember itself.[47] Most chronologies place the end of the Western Roman Empire in 476, when Romulus Augustulus was forced to abdicate to the Germanic warlord Odoacer.[48][49][50] Odoacer ended the Western Empire by declaring Zeno sole emperor and placing himself as Zeno's nominal subordinate. In reality, Italy was ruled by Odoacer alone.[48][49][51] The Eastern Roman Empire, called the Byzantine Empire by later historians, continued until the reign of Constantine XI Palaiologos, the last Roman emperor. He died in battle in 1453 against Mehmed II and his Ottoman forces during the siege of Constantinople. Mehmed II adopted the title of caesar in an attempt to claim a connection to the former Empire.[52][53] His claim was soon recognized by the Patriarchate of Constantinople, but not by European monarchs. ================================================ FILE: tmp/quantize_and_upload.py ================================================ #!/usr/bin/env python3 """ Download an mflux model, quantize it, and upload to HuggingFace. Usage (run from mflux project directory): cd /path/to/mflux uv run python /path/to/quantize_and_upload.py --model black-forest-labs/FLUX.1-Kontext-dev uv run python /path/to/quantize_and_upload.py --model black-forest-labs/FLUX.1-Kontext-dev --skip-base --skip-8bit uv run python /path/to/quantize_and_upload.py --model black-forest-labs/FLUX.1-Kontext-dev --dry-run Requires: - Must be run from mflux project directory using `uv run` - huggingface_hub installed (add to mflux deps or install separately) - HuggingFace authentication: run `huggingface-cli login` or set HF_TOKEN """ from __future__ import annotations import argparse import re import shutil import sys from pathlib import Path from typing import TYPE_CHECKING if TYPE_CHECKING: from mflux.models.flux.variants.txt2img.flux import Flux1 HF_ORG = "exolabs" def get_model_class(model_name: str) -> type: """Get the appropriate model class based on model name.""" from mflux.models.fibo.variants.txt2img.fibo import FIBO from mflux.models.flux.variants.txt2img.flux import Flux1 from mflux.models.flux2.variants.txt2img.flux2_klein import Flux2Klein from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage from mflux.models.z_image.variants.turbo.z_image_turbo import ZImageTurbo model_name_lower = model_name.lower() if "qwen" in model_name_lower: return QwenImage elif "fibo" in model_name_lower: return FIBO elif "z-image" in model_name_lower or "zimage" in model_name_lower: return ZImageTurbo elif "flux2" in model_name_lower or "flux.2" in model_name_lower: return Flux2Klein else: return Flux1 def get_repo_name(model_name: str, bits: int | None) -> str: """Get the HuggingFace repo name for a model variant.""" # Extract repo name from HF path (e.g., "black-forest-labs/FLUX.1-Kontext-dev" -> "FLUX.1-Kontext-dev") base_name = model_name.split("/")[-1] if "/" in model_name else model_name suffix = f"-{bits}bit" if bits else "" return f"{HF_ORG}/{base_name}{suffix}" def get_local_path(output_dir: Path, model_name: str, bits: int | None) -> Path: """Get the local save path for a model variant.""" # Extract repo name from HF path (e.g., "black-forest-labs/FLUX.1-Kontext-dev" -> "FLUX.1-Kontext-dev") base_name = model_name.split("/")[-1] if "/" in model_name else model_name suffix = f"-{bits}bit" if bits else "" return output_dir / f"{base_name}{suffix}" def copy_source_repo( source_repo: str, local_path: Path, dry_run: bool = False, ) -> None: """Copy all files from source repo (replicating original HF structure).""" print(f"\n{'=' * 60}") print(f"Copying full repo from source: {source_repo}") print(f"Output path: {local_path}") print(f"{'=' * 60}") if dry_run: print("[DRY RUN] Would download all files from source repo") return from huggingface_hub import snapshot_download # Download all files to our local path snapshot_download( repo_id=source_repo, local_dir=local_path, ) # Remove root-level safetensors files (flux.1-dev.safetensors, etc.) # These are redundant with the component directories for f in local_path.glob("*.safetensors"): print(f"Removing root-level safetensors: {f.name}") if not dry_run: f.unlink() print(f"Source repo copied to {local_path}") def load_and_save_quantized_model( model_name: str, bits: int, output_path: Path, dry_run: bool = False, ) -> None: """Load a model with quantization and save it in mflux format.""" print(f"\n{'=' * 60}") print(f"Loading {model_name} with {bits}-bit quantization...") print(f"Output path: {output_path}") print(f"{'=' * 60}") if dry_run: print("[DRY RUN] Would load and save quantized model") return from mflux.models.common.config.model_config import ModelConfig model_class = get_model_class(model_name) model_config = ModelConfig.from_name(model_name=model_name, base_model=None) model: Flux1 = model_class( quantize=bits, model_config=model_config, ) print(f"Saving model to {output_path}...") model.save_model(str(output_path)) print(f"Model saved successfully to {output_path}") def copy_source_metadata( source_repo: str, local_path: Path, dry_run: bool = False, ) -> None: """Copy metadata files (LICENSE, README, etc.) from source repo, excluding safetensors.""" print(f"\n{'=' * 60}") print(f"Copying metadata from source repo: {source_repo}") print(f"{'=' * 60}") if dry_run: print("[DRY RUN] Would download metadata files (excluding *.safetensors)") return from huggingface_hub import snapshot_download # Download all files except safetensors to our local path snapshot_download( repo_id=source_repo, local_dir=local_path, ignore_patterns=["*.safetensors"], ) print(f"Metadata files copied to {local_path}") def upload_to_huggingface( local_path: Path, repo_id: str, dry_run: bool = False, clean_remote: bool = False, ) -> None: """Upload a saved model to HuggingFace.""" print(f"\n{'=' * 60}") print(f"Uploading to HuggingFace: {repo_id}") print(f"Local path: {local_path}") print(f"Clean remote first: {clean_remote}") print(f"{'=' * 60}") if dry_run: print("[DRY RUN] Would upload to HuggingFace") return from huggingface_hub import HfApi api = HfApi() # Create the repo if it doesn't exist print(f"Creating/verifying repo: {repo_id}") api.create_repo(repo_id=repo_id, repo_type="model", exist_ok=True) # Clean remote repo if requested (delete old mflux-format files) if clean_remote: print("Cleaning old mflux-format files from remote...") try: # Pattern for mflux numbered shards: /.safetensors numbered_pattern = re.compile(r".*/\d+\.safetensors$") repo_files = api.list_repo_files(repo_id=repo_id, repo_type="model") for file_path in repo_files: # Delete numbered safetensors (mflux format) and mflux index files if numbered_pattern.match(file_path) or file_path.endswith( "/model.safetensors.index.json" ): print(f" Deleting: {file_path}") api.delete_file( path_in_repo=file_path, repo_id=repo_id, repo_type="model" ) except Exception as e: print(f"Warning: Could not clean remote files: {e}") # Upload the folder print("Uploading folder contents...") api.upload_folder( folder_path=str(local_path), repo_id=repo_id, repo_type="model", ) print(f"Upload complete: https://huggingface.co/{repo_id}") def clean_local_files(local_path: Path, dry_run: bool = False) -> None: """Remove local model files after upload.""" print(f"\nCleaning up: {local_path}") if dry_run: print("[DRY RUN] Would remove local files") return if local_path.exists(): shutil.rmtree(local_path) print(f"Removed {local_path}") def main() -> int: parser = argparse.ArgumentParser( description="Download an mflux model, quantize it, and upload to HuggingFace.", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: # Process all variants (base, 4-bit, 8-bit) for FLUX.1-Kontext-dev python tmp/quantize_and_upload.py --model black-forest-labs/FLUX.1-Kontext-dev # Only process 4-bit variant python tmp/quantize_and_upload.py --model black-forest-labs/FLUX.1-Kontext-dev --skip-base --skip-8bit # Save locally without uploading python tmp/quantize_and_upload.py --model black-forest-labs/FLUX.1-Kontext-dev --skip-upload # Preview what would happen python tmp/quantize_and_upload.py --model black-forest-labs/FLUX.1-Kontext-dev --dry-run """, ) parser.add_argument( "--model", "-m", required=True, help="HuggingFace model path (e.g., black-forest-labs/FLUX.1-Kontext-dev)", ) parser.add_argument( "--output-dir", type=Path, default=Path("./tmp/models"), help="Local directory to save models (default: ./tmp/models)", ) parser.add_argument( "--skip-base", action="store_true", help="Skip base model (no quantization)", ) parser.add_argument( "--skip-4bit", action="store_true", help="Skip 4-bit quantized model", ) parser.add_argument( "--skip-8bit", action="store_true", help="Skip 8-bit quantized model", ) parser.add_argument( "--skip-download", action="store_true", help="Skip downloading/processing, only do upload/clean operations", ) parser.add_argument( "--skip-upload", action="store_true", help="Only save locally, don't upload to HuggingFace", ) parser.add_argument( "--clean", action="store_true", help="Remove local files after upload", ) parser.add_argument( "--clean-remote", action="store_true", help="Delete old mflux-format files from remote repo before uploading", ) parser.add_argument( "--dry-run", action="store_true", help="Print actions without executing", ) args = parser.parse_args() # Determine which variants to process variants: list[int | None] = [] if not args.skip_base: variants.append(None) # Base model (no quantization) if not args.skip_4bit: variants.append(4) if not args.skip_8bit: variants.append(8) if not variants: print("Error: All variants skipped. Nothing to do.") return 1 # Create output directory args.output_dir.mkdir(parents=True, exist_ok=True) print(f"Model: {args.model}") print(f"Output directory: {args.output_dir}") print( f"Variants to process: {['base' if v is None else f'{v}-bit' for v in variants]}" ) print(f"Upload to HuggingFace: {not args.skip_upload}") print(f"Clean after upload: {args.clean}") if args.dry_run: print("\n*** DRY RUN MODE - No actual changes will be made ***") # Process each variant for bits in variants: local_path = get_local_path(args.output_dir, args.model, bits) repo_id = get_repo_name(args.model, bits) if not args.skip_download: if bits is None: # Base model: copy original HF repo structure (no mflux conversion) copy_source_repo( source_repo=args.model, local_path=local_path, dry_run=args.dry_run, ) else: # Quantized model: load, quantize, and save with mflux load_and_save_quantized_model( model_name=args.model, bits=bits, output_path=local_path, dry_run=args.dry_run, ) # Copy metadata from source repo (LICENSE, README, etc.) copy_source_metadata( source_repo=args.model, local_path=local_path, dry_run=args.dry_run, ) # Upload if not args.skip_upload: upload_to_huggingface( local_path=local_path, repo_id=repo_id, dry_run=args.dry_run, clean_remote=args.clean_remote, ) # Clean up if requested if args.clean: clean_local_files(local_path, dry_run=args.dry_run) print("\n" + "=" * 60) print("All done!") print("=" * 60) return 0 if __name__ == "__main__": sys.exit(main()) ================================================ FILE: tmp/run_llm.py ================================================ #!/usr/bin/env python3 import argparse import json import sys import requests def stream_chat(host: str, query: str) -> None: url = f"http://{host}:52415/v1/chat/completions" headers = {"Content-Type": "application/json"} payload = { "model": "mlx-community/Llama-3.2-1B-Instruct-4bit", # "model": "mlx-community/Llama-3_3-Nemotron-Super-49B-v1_5-mlx-4Bit", "stream": True, "messages": [{"role": "user", "content": query}], } try: with requests.post(url, headers=headers, json=payload, stream=True) as resp: resp.raise_for_status() for line in resp.iter_lines(decode_unicode=True): if not line: continue # SSE lines look like: "data: {...}" or "data: [DONE]" if not line.startswith("data:"): continue data = line[len("data:") :].strip() if data == "[DONE]": break try: obj = json.loads(data) except json.JSONDecodeError: continue for choice in obj.get("choices", []): delta = choice.get("delta") or {} content = delta.get("content") if content: print(content, end="", flush=True) except requests.RequestException as e: print(f"Request failed: {e}", file=sys.stderr) sys.exit(1) print() def main() -> None: parser = argparse.ArgumentParser( description="Stream chat completions from a local server." ) parser.add_argument("host", help="Hostname (without protocol), e.g. localhost") parser.add_argument( "-f", "--file", help="Path to a text file whose contents will be used as the query", ) parser.add_argument( "query", nargs="*", help="Query text (if not using -f/--file). All remaining arguments are joined with spaces.", ) args = parser.parse_args() if args.file: try: with open(args.file, "r", encoding="utf-8") as f: query = f.read().strip() except OSError as e: print(f"Error reading file {args.file}: {e}", file=sys.stderr) sys.exit(1) elif args.query: query = " ".join(args.query) else: parser.error("You must provide either a query or a file (-f/--file).") stream_chat(args.host, query) if __name__ == "__main__": main() ================================================ FILE: tmp/run_llm.sh ================================================ #!/usr/bin/env bash set -euo pipefail if [ $# -lt 2 ]; then echo "Usage: $0 " exit 1 fi HOST="$1" shift QUERY="$*" curl -sN -X POST "http://$HOST:52415/v1/chat/completions" \ -H "Content-Type: application/json" \ -d "{ \"model\": \"mlx-community/Kimi-K2-Thinking\", \"stream\": true, \"messages\": [{ \"role\": \"user\", \"content\": \"$QUERY\"}] }" | grep --line-buffered '^data:' | grep --line-buffered -v 'data: \[DONE\]' | cut -d' ' -f2- | jq -r --unbuffered '.choices[].delta.content // empty' | awk '{ORS=""; print; fflush()} END {print "\n"}' ================================================ FILE: tmp/set_rdma_network_config.sh ================================================ #!/usr/bin/env bash set -euo pipefail PREFS="/Library/Preferences/SystemConfiguration/preferences.plist" # Remove bridge0 interface ifconfig bridge0 &>/dev/null && { ifconfig bridge0 | grep -q 'member' && { ifconfig bridge0 | awk '/member/ {print $2}' | xargs -n1 ifconfig bridge0 deletem 2>/dev/null || true } ifconfig bridge0 destroy 2>/dev/null || true } # Remove Thunderbolt Bridge from VirtualNetworkInterfaces in preferences.plist /usr/libexec/PlistBuddy -c "Delete :VirtualNetworkInterfaces:Bridge:bridge0" "$PREFS" 2>/dev/null || true networksetup -listlocations | grep -q exo || { networksetup -createlocation exo } networksetup -switchtolocation exo networksetup -listallhardwareports | awk -F': ' '/Hardware Port: / {print $2}' | while IFS=":" read -r name; do case "$name" in "Ethernet Adapter"*) ;; "Thunderbolt Bridge") ;; "Thunderbolt "*) networksetup -listallnetworkservices | grep -q "EXO $name" || networksetup -createnetworkservice "EXO $name" "$name" 2>/dev/null || continue networksetup -setdhcp "EXO $name" ;; *) networksetup -listallnetworkservices | grep -q "$name" || networksetup -createnetworkservice "$name" "$name" 2>/dev/null || continue ;; esac done networksetup -listnetworkservices | grep -q "Thunderbolt Bridge" && { networksetup -setnetworkserviceenabled "Thunderbolt Bridge" off } || true ================================================ FILE: tmp/test_trust_remote_code_attack.sh ================================================ #!/usr/bin/env bash # Test that models added via API get trust_remote_code=false # Run this against a running exo instance. # Usage: ./test_trust_remote_code_attack.sh [host:port] set -uo pipefail HOST="${1:-localhost:52415}" MODEL_ID="KevTheHermit/security-testing" CUSTOM_CARDS_DIR="$HOME/.exo/custom_model_cards" CARD_FILE="$CUSTOM_CARDS_DIR/KevTheHermit--security-testing.toml" echo "=== Test: trust_remote_code attack via API ===" echo "Target: $HOST" echo "" # Clean up RCE proof from previous runs rm -f /tmp/exo-rce-proof.txt # Step 0: Clean up any stale card from previous runs if [ -f "$CARD_FILE" ]; then echo "[0] Removing stale card from previous run ..." curl -s -X DELETE \ "http://$HOST/models/custom/$(python3 -c 'import urllib.parse; print(urllib.parse.quote("'"$MODEL_ID"'", safe=""))')" >/dev/null rm -f "$CARD_FILE" echo " Done" echo "" fi # Step 1: Add the malicious model via API echo "[1] Adding model via POST /models/add ..." ADD_RESPONSE=$(curl -s -w "\n%{http_code}" -X POST "http://$HOST/models/add" \ -H "Content-Type: application/json" \ -d "{\"model_id\":\"$MODEL_ID\"}") HTTP_CODE=$(echo "$ADD_RESPONSE" | tail -1) BODY=$(echo "$ADD_RESPONSE" | sed '$d') echo " HTTP $HTTP_CODE" if [ "$HTTP_CODE" -ge 400 ]; then echo " Model add failed (HTTP $HTTP_CODE) — that's fine if model doesn't exist on HF." echo " Response: $BODY" echo "" echo "RESULT: Model was rejected at add time. Attack blocked." exit 0 fi # Step 2: Verify the saved TOML has trust_remote_code = false echo "" echo "[2] Checking saved model card TOML ..." if [ ! -f "$CARD_FILE" ]; then echo " FAIL: Card file not found at $CARD_FILE" exit 1 fi if grep -q 'trust_remote_code = false' "$CARD_FILE"; then echo " SAFE: trust_remote_code = false (fix is active)" else echo " VULNERABLE: trust_remote_code is not false — remote code WILL be trusted" fi echo " Contents:" cat "$CARD_FILE" # Step 3: Place the instance echo "" echo "[3] Attempting POST /place_instance ..." PLACE_RESPONSE=$(curl -s -w "\n%{http_code}" -X POST "http://$HOST/place_instance" \ -H "Content-Type: application/json" \ -d "{\"model_id\":\"$MODEL_ID\"}") PLACE_CODE=$(echo "$PLACE_RESPONSE" | tail -1) PLACE_BODY=$(echo "$PLACE_RESPONSE" | sed '$d') echo " HTTP $PLACE_CODE" echo " Response: $PLACE_BODY" # Step 3b: Send a chat completion to actually trigger tokenizer loading echo "" echo "[3b] Sending chat completion to trigger tokenizer load ..." CHAT_RESPONSE=$(curl -s -w "\n%{http_code}" --max-time 30 -X POST "http://$HOST/v1/chat/completions" \ -H "Content-Type: application/json" \ -d "{\"model\":\"$MODEL_ID\",\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}],\"max_tokens\":1}") CHAT_CODE=$(echo "$CHAT_RESPONSE" | tail -1) CHAT_BODY=$(echo "$CHAT_RESPONSE" | sed '$d') echo " HTTP $CHAT_CODE" echo " Response: $CHAT_BODY" echo "" echo "[3c] Checking for RCE proof ..." sleep 5 if [ -f /tmp/exo-rce-proof.txt ]; then echo " VULNERABLE: Remote code executed!" echo " Contents:" cat /tmp/exo-rce-proof.txt else echo " SAFE: /tmp/exo-rce-proof.txt does not exist — remote code was NOT executed" fi # Step 4: Clean up — delete instance and custom model echo "" echo "[4] Cleaning up ..." # Find and delete any instance for this model INSTANCE_ID=$(curl -s "http://$HOST/state" | python3 -c " import sys, json state = json.load(sys.stdin) for iid, wrapper in state.get('instances', {}).items(): for tag, inst in wrapper.items(): sa = inst.get('shardAssignments', {}) if sa.get('modelId', '') == '$MODEL_ID': print(iid) sys.exit(0) " 2>/dev/null || true) if [ -n "$INSTANCE_ID" ]; then echo " Deleting instance $INSTANCE_ID ..." curl -s -X DELETE "http://$HOST/instance/$INSTANCE_ID" >/dev/null echo " Done" else echo " No instance found to delete" fi echo " Deleting custom model card ..." curl -s -X DELETE \ "http://$HOST/models/custom/$(python3 -c 'import urllib.parse; print(urllib.parse.quote("'"$MODEL_ID"'", safe=""))')" >/dev/null echo " Done" echo "" echo "=== DONE ==="