Adding trust_remote_code to the orchestrator and worker

This commit is contained in:
samit 2026-03-08 16:44:41 -07:00
commit 662cb1c440
2 changed files with 4 additions and 0 deletions

View file

@ -211,6 +211,7 @@ class ExportOrchestrator:
checkpoint_path: str,
max_seq_length: int = 2048,
load_in_4bit: bool = True,
trust_remote_code: bool = False,
) -> Tuple[bool, str]:
"""Load a checkpoint for export.
@ -225,6 +226,7 @@ class ExportOrchestrator:
"checkpoint_path": checkpoint_path,
"max_seq_length": max_seq_length,
"load_in_4bit": load_in_4bit,
"trust_remote_code": trust_remote_code,
}
# Always kill existing subprocess and spawn fresh.

View file

@ -84,6 +84,7 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
checkpoint_path = cmd["checkpoint_path"]
max_seq_length = cmd.get("max_seq_length", 2048)
load_in_4bit = cmd.get("load_in_4bit", True)
trust_remote_code = cmd.get("trust_remote_code", False)
try:
_send_response(resp_queue, {
@ -96,6 +97,7 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
checkpoint_path=checkpoint_path,
max_seq_length=max_seq_length,
load_in_4bit=load_in_4bit,
trust_remote_code=trust_remote_code,
)
_send_response(resp_queue, {