tmp changes - DONT PUSH

src/exo/main.py

@@ -261,13 +261,6 @@ def main():
     if args.offline:
         logger.info("Running in OFFLINE mode — no internet checks, local models only")
 
-    # Set trust_remote_code override env var for runner subprocesses
-    if args.trust_remote_code:
-        os.environ["EXO_TRUST_REMOTE_CODE"] = "1"
-        logger.warning(
-            "--trust-remote-code enabled: models may execute arbitrary code during loading"
-        )
-
     # Set FAST_SYNCH override env var for runner subprocesses
     if args.fast_synch is True:
         os.environ["EXO_FAST_SYNCH"] = "on"
@@ -292,7 +285,6 @@ class Args(CamelCaseModel):
     no_downloads: bool = False
     offline: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
-    trust_remote_code: bool = False
 
     @classmethod
     def parse(cls) -> Self:
@@ -344,11 +336,6 @@ class Args(CamelCaseModel):
             action="store_true",
             help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models",
         )
-        parser.add_argument(
-            "--trust-remote-code",
-            action="store_true",
-            help="Allow models to execute custom code during tokenizer loading (security-sensitive, CLI-only)",
-        )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(
             "--fast-synch",

src/exo/master/tests/test_placement.py

@@ -14,12 +14,10 @@ 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.events import InstanceCreated, InstanceDeleted
 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,
@@ -458,117 +456,3 @@ def test_tensor_rdma_backend_connectivity_matrix(
         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

src/exo/worker/engines/mlx/constants.py

@@ -13,6 +13,5 @@ KV_CACHE_BITS: int | None = None
 
 DEFAULT_TOP_LOGPROBS: int = 5
 
-# True for built-in models with known model cards; custom models added via API default to False
-# and can be overridden with the --trust-remote-code CLI flag.
+# TODO: We should really make this opt-in, but Kimi requires trust_remote_code=True
 TRUST_REMOTE_CODE: bool = True

src/exo/worker/engines/mlx/utils_mlx.py

@@ -291,14 +291,10 @@ def shard_and_load(
 
 def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata) -> TokenizerWrapper:
     """Load tokenizer for a model shard. Delegates to load_tokenizer_for_model_id."""
-    trust_remote_code = (
-        shard_metadata.model_card.trust_remote_code
-        or os.environ.get("EXO_TRUST_REMOTE_CODE") == "1"
-    )
     return load_tokenizer_for_model_id(
         shard_metadata.model_card.model_id,
         model_path,
-        trust_remote_code=trust_remote_code,
+        trust_remote_code=shard_metadata.model_card.trust_remote_code,
     )
 
 

tests/get_all_models_on_cluster.py

@@ -8,7 +8,7 @@ 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
+    ["/Applications/Tailscale.app/Contents/MacOS/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

tests/start_distributed_test.py

@@ -15,7 +15,7 @@ if not (args := sys.argv[1:]):
 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
+    ["/Applications/Tailscale.app/Contents/MacOS/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]