@tuned-tensor/local 0.2.6 → 0.2.8

package/training/local-runner/src/prefetch.py

@@ -0,0 +1,178 @@
+import argparse
+import hashlib
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+ALLOW_PATTERNS = [
+    "*.json",
+    "*.model",
+    "*.py",
+    "*.safetensors",
+    "*.safetensors.index.json",
+    "*.tiktoken",
+    "*.txt",
+    "chat_template.jinja",
+    "merges.txt",
+    "pytorch_model*.bin",
+    "sentencepiece.bpe.model",
+    "spiece.model",
+    "tokenizer.*",
+    "vocab.*",
+]
+
+IGNORE_PATTERNS = [
+    "*.ckpt",
+    "*.gguf",
+    "*.h5",
+    "*.msgpack",
+    "*.onnx",
+    "*.ot",
+    "optimizer.pt",
+    "rng_state*.pth",
+    "scheduler.pt",
+    "trainer_state.json",
+    "training_args.bin",
+]
+
+
+def write_json(path: str, value: dict[str, Any]) -> None:
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    destination.write_text(json.dumps(value, indent=2) + "\n", encoding="utf-8")
+
+
+def configure_hugging_face_cache(cache_home: str | None) -> None:
+    """Treat model_cache as HF_HOME and keep legacy overrides consistent."""
+    if not cache_home:
+        return
+    # Keep the lexical absolute path supplied by the Node runner. Path.resolve
+    # canonicalizes macOS /var to /private/var, making metadata disagree with
+    # training even though both locations identify the same directory.
+    home = Path(cache_home).expanduser().absolute()
+    hub = home / "hub"
+    os.environ.update({
+        "HF_HOME": str(home),
+        "HF_HUB_CACHE": str(hub),
+        "HUGGINGFACE_HUB_CACHE": str(hub),
+    })
+    for deprecated in (
+        "TRANSFORMERS_CACHE",
+        "PYTORCH_TRANSFORMERS_CACHE",
+        "PYTORCH_PRETRAINED_BERT_CACHE",
+    ):
+        os.environ.pop(deprecated, None)
+
+
+def hash_file(path: Path, algorithm: str, prefix: bytes = b"") -> str:
+    digest = hashlib.new(algorithm)
+    digest.update(prefix)
+    with path.open("rb") as source:
+        for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def verify_snapshot(snapshot: Path) -> tuple[list[Path], int]:
+    config = snapshot / "config.json"
+    if not config.is_file() or config.stat().st_size == 0:
+        raise ValueError(f"Cached snapshot is missing a non-empty config.json: {snapshot}")
+    try:
+        json.loads(config.read_text(encoding="utf-8"))
+    except (OSError, json.JSONDecodeError) as exc:
+        raise ValueError(f"Cached snapshot has an invalid config.json: {snapshot}") from exc
+
+    weight_files = sorted(snapshot.glob("*.safetensors")) + sorted(snapshot.glob("pytorch_model*.bin"))
+    index_files = sorted(snapshot.glob("*.safetensors.index.json")) + sorted(snapshot.glob("pytorch_model*.bin.index.json"))
+    for index in index_files:
+        try:
+            weight_map = json.loads(index.read_text(encoding="utf-8")).get("weight_map", {})
+        except (OSError, json.JSONDecodeError) as exc:
+            raise ValueError(f"Cached snapshot has an invalid weight index: {index}") from exc
+        if not isinstance(weight_map, dict) or not weight_map:
+            raise ValueError(f"Cached snapshot has an empty weight index: {index}")
+        for name in set(weight_map.values()):
+            if not isinstance(name, str) or not (snapshot / name).is_file():
+                raise ValueError(f"Cached snapshot is missing an indexed weight shard: {name}")
+    if not weight_files:
+        raise ValueError(f"Cached snapshot contains no Transformers model weights: {snapshot}")
+    for weight in weight_files:
+        if weight.stat().st_size == 0:
+            raise ValueError(f"Cached snapshot contains an empty model weight: {weight}")
+
+    tokenizer_config = snapshot / "tokenizer_config.json"
+    vocabulary = [
+        snapshot / "tokenizer.json",
+        snapshot / "tokenizer.model",
+        snapshot / "sentencepiece.bpe.model",
+        snapshot / "spiece.model",
+        snapshot / "vocab.json",
+        snapshot / "tokenizer.tiktoken",
+    ]
+    if not tokenizer_config.is_file() or not any(path.is_file() and path.stat().st_size > 0 for path in vocabulary):
+        raise ValueError(f"Cached snapshot is missing tokenizer metadata or vocabulary files: {snapshot}")
+
+    files = [path for path in snapshot.rglob("*") if path.is_file()]
+    verified_blobs = 0
+    for path in files:
+        blob = path.resolve()
+        expected = blob.name.lower()
+        if len(expected) == 64 and all(character in "0123456789abcdef" for character in expected):
+            actual = hash_file(blob, "sha256")
+        elif len(expected) == 40 and all(character in "0123456789abcdef" for character in expected):
+            actual = hash_file(blob, "sha1", f"blob {blob.stat().st_size}\0".encode())
+        else:
+            continue
+        if actual != expected:
+            raise ValueError(f"Cached Hugging Face blob checksum mismatch: {path}")
+        verified_blobs += 1
+    return files, verified_blobs
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Prefetch a TT Local Hugging Face base model.")
+    parser.add_argument("--input", required=True)
+    parser.add_argument("--output", required=True)
+    args = parser.parse_args()
+
+    payload = json.loads(Path(args.input).read_text(encoding="utf-8"))
+    base_model = str(payload["base_model"])
+    cache_dir = payload.get("model_cache")
+    configure_hugging_face_cache(cache_dir)
+
+    # Import only after the cache environment is final. huggingface_hub reads
+    # these variables while importing its constants module.
+    from huggingface_hub import constants, snapshot_download
+
+    action = "Verifying cached" if payload.get("local_files_only") else "Prefetching"
+    print(f"{action} {base_model} in {constants.HF_HUB_CACHE}...", flush=True)
+    snapshot_path = snapshot_download(
+        repo_id=base_model,
+        revision=payload.get("revision"),
+        token=os.getenv("HF_TOKEN"),
+        allow_patterns=ALLOW_PATTERNS,
+        ignore_patterns=IGNORE_PATTERNS,
+        local_files_only=bool(payload.get("local_files_only", False)),
+    )
+    snapshot = Path(snapshot_path)
+    print(f"Verifying snapshot files under {snapshot}...", flush=True)
+    snapshot_files, verified_blob_count = verify_snapshot(snapshot)
+
+    write_json(args.output, {
+        "ok": True,
+        "base_model": base_model,
+        "loader": payload.get("loader"),
+        "model_cache": str(constants.HF_HOME),
+        "hf_home": str(constants.HF_HOME),
+        "hub_cache": str(constants.HF_HUB_CACHE),
+        "snapshot_path": snapshot_path,
+        "snapshot_revision": Path(snapshot_path).name,
+        "file_count": len(snapshot_files),
+        "size_bytes": sum(path.stat().st_size for path in snapshot_files),
+        "verified_blob_count": verified_blob_count,
+    })
+
+
+if __name__ == "__main__":
+    main()

package/training/local-runner/src/serve.py

@@ -0,0 +1,329 @@
+from __future__ import annotations
+
+import json
+import base64
+import hmac
+import math
+import os
+import sys
+import threading
+import time
+import uuid
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any
+
+# HF_HOME is supplied by the Node launcher before this module starts. Importing
+# the model stack only after reading process configuration keeps every local
+# workflow on the same Hugging Face cache contract.
+from evaluate import (  # noqa: E402
+    configure_hugging_face_cache,
+    generate_multimodal_one,
+    import_runtime_dependencies,
+    load_multimodal_model,
+    load_text_model,
+    resolve_adapter_path,
+    sampling_kwargs,
+)
+
+
+MODEL_ARTIFACT = os.environ["TT_MODEL_ARTIFACT"]
+BASE_MODEL = os.environ["TT_BASE_MODEL"]
+BASE_MODEL_REVISION = os.environ.get("TT_BASE_MODEL_REVISION")
+MODEL_NAME = os.environ.get("TT_MODEL_NAME", "tuned-tensor-local")
+MODEL_LOADER = os.environ.get("TT_MODEL_LOADER", "causal_lm")
+SYSTEM_PROMPT = os.environ.get("TT_SYSTEM_PROMPT", "").strip()
+HOST = os.environ.get("TT_HOST", "127.0.0.1")
+PORT = int(os.environ.get("TT_PORT", "8000"))
+DEVICE_REQUEST = os.environ.get("TT_DEVICE", "auto")
+TRUST_REMOTE_CODE = os.environ.get("TT_TRUST_REMOTE_CODE", "true").lower() in {"1", "true", "yes"}
+DEFAULT_MAX_TOKENS = int(os.environ.get("TT_MAX_TOKENS", "512"))
+DEFAULT_TEMPERATURE = float(os.environ.get("TT_TEMPERATURE", "0"))
+DEFAULT_TOP_P = float(os.environ.get("TT_TOP_P", "1"))
+CHAT_TEMPLATE_KWARGS = json.loads(os.environ.get("TT_CHAT_TEMPLATE_KWARGS", "{}"))
+MAX_REQUEST_BYTES = 2 * 1024 * 1024
+MAX_PROMPT_CHARS = 100_000
+MAX_PROMPT_TOKENS = 16_384
+MAX_MESSAGES = 128
+MAX_IMAGE_BYTES = 5 * 1024 * 1024
+API_KEY = os.environ.get("TT_API_KEY", "")
+MAX_CONCURRENT_REQUESTS = int(os.environ.get("TT_MAX_CONCURRENT_REQUESTS", "1"))
+
+
+configure_hugging_face_cache(os.environ.get("HF_HOME"))
+import_runtime_dependencies()
+TEMP_DIR = TemporaryDirectory(prefix="tt-local-serve-")
+ADAPTER_PATH = resolve_adapter_path(MODEL_ARTIFACT, Path(TEMP_DIR.name))
+MODEL_PAYLOAD = {
+    "base_model": BASE_MODEL,
+    "base_model_revision": BASE_MODEL_REVISION,
+    "device": DEVICE_REQUEST,
+    "trust_remote_code": TRUST_REMOTE_CODE,
+    "model_loader": MODEL_LOADER,
+}
+
+if MODEL_LOADER == "image_text_to_text":
+    MODEL, PROCESSOR, DEVICE = load_multimodal_model(MODEL_PAYLOAD, ADAPTER_PATH)
+    TOKENIZER = getattr(PROCESSOR, "tokenizer", PROCESSOR)
+else:
+    MODEL, TOKENIZER, DEVICE = load_text_model(MODEL_PAYLOAD, ADAPTER_PATH)
+    PROCESSOR = None
+
+GENERATION_LOCK = threading.Lock()
+REQUEST_SLOTS = threading.BoundedSemaphore(MAX_CONCURRENT_REQUESTS)
+
+
+def normalize_messages(raw_messages: Any) -> list[dict[str, Any]]:
+    if not isinstance(raw_messages, list) or not raw_messages:
+        raise ValueError("Request must include a non-empty messages array.")
+    if len(raw_messages) > MAX_MESSAGES:
+        raise ValueError(f"Request exceeds the {MAX_MESSAGES}-message limit.")
+    messages: list[dict[str, Any]] = []
+    for raw in raw_messages:
+        if not isinstance(raw, dict) or raw.get("role") not in {"system", "user", "assistant", "tool"}:
+            raise ValueError("Each message must have a supported role and content.")
+        messages.append({"role": raw["role"], "content": raw.get("content", "")})
+    # A spec passed by the server owner is invariant behavior. Client-supplied
+    # system messages may add context but cannot replace it.
+    if SYSTEM_PROMPT:
+        messages.insert(0, {"role": "system", "content": SYSTEM_PROMPT})
+    if sum(len(text_content(message.get("content"))) for message in messages) > MAX_PROMPT_CHARS:
+        raise ValueError(f"Prompt exceeds the {MAX_PROMPT_CHARS}-character limit.")
+    return messages
+
+
+def text_content(content: Any) -> str:
+    if isinstance(content, str):
+        return content
+    if not isinstance(content, list):
+        return str(content) if content is not None else ""
+    parts: list[str] = []
+    for part in content:
+        if isinstance(part, dict) and part.get("type") == "text":
+            parts.append(str(part.get("text", "")))
+        elif isinstance(part, dict) and part.get("type") in {"image", "image_url"}:
+            continue
+        else:
+            parts.append(str(part))
+    return "\n".join(part for part in parts if part)
+
+
+def image_assets(messages: list[dict[str, Any]]) -> list[dict[str, str]]:
+    assets: list[dict[str, str]] = []
+    for message in messages:
+        content = message.get("content")
+        if not isinstance(content, list):
+            continue
+        for part in content:
+            if not isinstance(part, dict) or part.get("type") not in {"image", "image_url"}:
+                continue
+            image = part.get("image_url")
+            if isinstance(image, dict):
+                image = image.get("url")
+            image = image or part.get("image") or part.get("data_uri") or part.get("uri") or part.get("path")
+            if isinstance(image, str) and image:
+                if not image.startswith("data:image/") or "," not in image:
+                    raise ValueError("Serving accepts image content only as data:image/... URIs.")
+                encoded = image.split(",", 1)[1]
+                try:
+                    decoded_size = len(base64.b64decode(encoded, validate=True))
+                except Exception as exc:
+                    raise ValueError("Image data URI contains invalid base64.") from exc
+                if decoded_size > MAX_IMAGE_BYTES:
+                    raise ValueError(f"Image exceeds the {MAX_IMAGE_BYTES}-byte decoded limit.")
+                assets.append({"type": "image", "image": image})
+    return assets
+
+
+def generate_text(messages: list[dict[str, Any]], generation: dict[str, Any]) -> tuple[str, int, int]:
+    normalized = [{"role": message["role"], "content": text_content(message.get("content"))} for message in messages]
+    try:
+        prompt = TOKENIZER.apply_chat_template(
+            normalized,
+            tokenize=False,
+            add_generation_prompt=True,
+            **CHAT_TEMPLATE_KWARGS,
+        )
+    except Exception:
+        prompt = "\n".join(f"{message['role']}: {message['content']}" for message in normalized) + "\nassistant:"
+    inputs = TOKENIZER(prompt, return_tensors="pt")
+    if int(inputs["input_ids"].shape[-1]) > MAX_PROMPT_TOKENS:
+        raise ValueError(f"Prompt exceeds the {MAX_PROMPT_TOKENS}-token limit.")
+    target_device = next(MODEL.parameters()).device
+    inputs = {key: value.to(target_device) for key, value in inputs.items()}
+    with GENERATION_LOCK:
+        import torch
+
+        with torch.no_grad():
+            generated = MODEL.generate(
+                **inputs,
+                max_new_tokens=int(generation["max_new_tokens"]),
+                pad_token_id=TOKENIZER.eos_token_id,
+                **sampling_kwargs(generation),
+            )
+    prompt_tokens = int(inputs["input_ids"].shape[-1])
+    completion_tokens = generated[0][prompt_tokens:]
+    content = TOKENIZER.decode(completion_tokens, skip_special_tokens=True).strip()
+    return content, prompt_tokens, int(completion_tokens.shape[-1])
+
+
+def generate_multimodal(messages: list[dict[str, Any]], generation: dict[str, Any]) -> tuple[str, int, int]:
+    system_parts = [text_content(message.get("content")) for message in messages if message["role"] == "system"]
+    conversation = "\n".join(
+        f"{message['role']}: {text_content(message.get('content'))}"
+        for message in messages
+        if message["role"] != "system"
+    )
+    example = {
+        "input": conversation,
+        "output": "",
+        "input_assets": image_assets(messages),
+    }
+    with GENERATION_LOCK:
+        result = generate_multimodal_one(
+            MODEL,
+            PROCESSOR,
+            "\n\n".join(part for part in system_parts if part),
+            example,
+            generation,
+            Path.cwd(),
+            CHAT_TEMPLATE_KWARGS,
+        )
+    # The evaluator does not currently expose token counts for multimodal
+    # requests. Returning zero is preferable to inventing usage data.
+    return str(result["actual"]), 0, 0
+
+
+def bounded_number(value: Any, default: float, minimum: float, maximum: float) -> float:
+    number = float(default if value is None else value)
+    if not math.isfinite(number) or number < minimum or number > maximum:
+        raise ValueError(f"Generation value must be between {minimum} and {maximum}.")
+    return number
+
+
+def bounded_integer(value: Any, default: int, minimum: int, maximum: int) -> int:
+    number = bounded_number(value, default, minimum, maximum)
+    if not number.is_integer():
+        raise ValueError(f"Generation value must be an integer between {minimum} and {maximum}.")
+    return int(number)
+
+
+class Handler(BaseHTTPRequestHandler):
+    server_version = "TunedTensorLocalServer/0.1"
+
+    def setup(self) -> None:
+        super().setup()
+        self.connection.settimeout(30)
+
+    def authorized(self) -> bool:
+        if not API_KEY:
+            return True
+        supplied = self.headers.get("authorization", "")
+        expected = "Bearer " + API_KEY
+        return hmac.compare_digest(supplied, expected)
+
+    def send_json(self, status: int, payload: Any) -> None:
+        body = json.dumps(payload).encode("utf-8")
+        self.send_response(status)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def do_GET(self) -> None:  # noqa: N802
+        if not self.authorized():
+            self.send_json(401, {"error": {"message": "Unauthorized"}})
+            return
+        if self.path == "/health":
+            self.send_json(200, {"status": "ok", "model": MODEL_NAME, "device": DEVICE})
+            return
+        if self.path == "/v1/models":
+            self.send_json(200, {
+                "object": "list",
+                "data": [{"id": MODEL_NAME, "object": "model", "owned_by": "tuned-tensor-local"}],
+            })
+            return
+        self.send_json(404, {"error": {"message": "Not found"}})
+
+    def do_POST(self) -> None:  # noqa: N802
+        if not self.authorized():
+            self.send_json(401, {"error": {"message": "Unauthorized"}})
+            return
+        if self.path not in {"/v1/chat/completions", "/chat/completions"}:
+            self.send_json(404, {"error": {"message": "Not found"}})
+            return
+        try:
+            length = int(self.headers.get("content-length", "0"))
+            if length <= 0 or length > MAX_REQUEST_BYTES:
+                raise ValueError("Request body is empty or too large.")
+            body = json.loads(self.rfile.read(length))
+            if not isinstance(body, dict):
+                raise ValueError("Request body must be a JSON object.")
+            if body.get("stream"):
+                self.send_json(400, {"error": {"message": "Streaming is not supported yet."}})
+                return
+            messages = normalize_messages(body.get("messages"))
+            generation = {
+                "max_new_tokens": bounded_integer(body.get("max_tokens"), DEFAULT_MAX_TOKENS, 1, 8192),
+                "temperature": bounded_number(body.get("temperature"), DEFAULT_TEMPERATURE, 0, 5),
+                "top_p": bounded_number(body.get("top_p"), DEFAULT_TOP_P, 0, 1),
+            }
+            if not REQUEST_SLOTS.acquire(blocking=False):
+                self.send_json(429, {"error": {"message": "The local model is busy; retry shortly."}})
+                return
+            started = time.perf_counter()
+            try:
+                if MODEL_LOADER == "image_text_to_text":
+                    content, prompt_tokens, completion_tokens = generate_multimodal(messages, generation)
+                else:
+                    if image_assets(messages):
+                        raise ValueError("This text model does not accept image content parts.")
+                    content, prompt_tokens, completion_tokens = generate_text(messages, generation)
+            finally:
+                REQUEST_SLOTS.release()
+            latency_ms = round((time.perf_counter() - started) * 1000)
+            self.send_json(200, {
+                "id": "chatcmpl-" + uuid.uuid4().hex,
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": MODEL_NAME,
+                "choices": [{
+                    "index": 0,
+                    "message": {"role": "assistant", "content": content},
+                    "finish_reason": "stop",
+                }],
+                "usage": {
+                    "prompt_tokens": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": prompt_tokens + completion_tokens,
+                },
+                "tt_local": {"latency_ms": latency_ms, "device": DEVICE},
+            })
+        except (ValueError, TypeError, json.JSONDecodeError) as exc:
+            self.send_json(400, {"error": {"message": str(exc)}})
+        except Exception as exc:
+            print(f"Model server error: {type(exc).__name__}: {exc}", file=sys.stderr, flush=True)
+            self.send_json(500, {"error": {"message": "Internal model server error."}})
+
+    def log_message(self, fmt: str, *args: Any) -> None:
+        print("%s - %s" % (self.address_string(), fmt % args), flush=True)
+
+
+def main() -> None:
+    print(f"Serving {MODEL_NAME} on http://{HOST}:{PORT}", flush=True)
+    print(f"OpenAI-compatible endpoint: http://{HOST}:{PORT}/v1/chat/completions", flush=True)
+    print(f"Device: {DEVICE}", flush=True)
+    server = ThreadingHTTPServer((HOST, PORT), Handler)
+    server.daemon_threads = True
+    server.request_queue_size = max(2, MAX_CONCURRENT_REQUESTS * 2)
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        print("Model server stopped.", flush=True)
+    finally:
+        server.server_close()
+
+
+if __name__ == "__main__":
+    main()

package/training/local-runner/src/train.py

@@ -36,6 +36,8 @@ HYPERPARAMETERS_PATH = Path(
 MODEL_DIR = Path(os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
 OUTPUT_DIR = Path(os.environ.get("SM_OUTPUT_DIR", "/opt/ml/output"))
 ROW_SOURCE_DIR_KEY = "_tt_source_dir"
+MAX_ARCHIVE_MEMBERS = 20_000
+MAX_ARCHIVE_EXPANDED_BYTES = 20 * 1024 * 1024 * 1024
 
 
 def load_hyperparameters() -> dict[str, str]:
@@ -70,6 +72,11 @@ def hp_bool(name: str, default: bool) -> bool:
     return (hp(name, str(default)) or "").lower() in {"1", "true", "yes", "y"}
 
 
+def model_revision_kwargs(model_source: str) -> dict[str, str]:
+    revision = hp("base_model_revision")
+    return {"revision": revision} if revision and not Path(model_source).exists() else {}
+
+
 def load_rows() -> list[dict[str, Any]]:
     rows: list[dict[str, Any]] = []
     for path in sorted(TRAINING_DIR.rglob("*.jsonl")):
@@ -86,15 +93,36 @@ def load_rows() -> list[dict[str, Any]]:
     return rows
 
 
-def resolve_model_source() -> str:
+def safe_extract_archive(path: Path, destination: Path) -> None:
+    destination.mkdir(parents=True, exist_ok=True)
+    destination_root = destination.resolve()
+    with tarfile.open(path, "r:gz") as tar:
+        members = tar.getmembers()
+        if len(members) > MAX_ARCHIVE_MEMBERS:
+            raise ValueError(f"Model archive exceeds {MAX_ARCHIVE_MEMBERS} members")
+        expanded_bytes = sum(max(0, member.size) for member in members if member.isfile())
+        if expanded_bytes > MAX_ARCHIVE_EXPANDED_BYTES:
+            raise ValueError("Model archive exceeds the 20 GiB expanded-size limit")
+        for member in members:
+            destination_path = (destination / member.name).resolve()
+            try:
+                destination_path.relative_to(destination_root)
+            except ValueError:
+                raise ValueError(f"Unsafe archive member: {member.name}")
+            if member.issym() or member.islnk() or member.isdev():
+                raise ValueError(f"Unsafe archive member type: {member.name}")
+        tar.extractall(destination)
+
+
+def resolve_model_source(tmp: Path) -> str:
     archive = BASE_MODEL_DIR / "model.tar.gz"
     if archive.is_file():
-        extracted = Path("/tmp/base-model")
-        extracted.mkdir(parents=True, exist_ok=True)
-        with tarfile.open(archive, "r:gz") as tar:
-            tar.extractall(extracted)
+        extracted = tmp / "base-model"
+        safe_extract_archive(archive, extracted)
         candidates = [path for path in extracted.iterdir() if path.is_dir()]
         return str(candidates[0] if len(candidates) == 1 else extracted)
+    if BASE_MODEL_DIR.is_dir() and any(BASE_MODEL_DIR.iterdir()):
+        return str(BASE_MODEL_DIR)
     return hp("base_model") or "Qwen/Qwen3.5-2B"
 
 
@@ -113,16 +141,7 @@ def resolve_adapter_path(value: str | None, tmp: Path) -> str | None:
     path = Path(path_value)
     if path.is_file() and path.name.endswith(".tar.gz"):
         extracted = tmp / "parent-adapter"
-        extracted.mkdir(parents=True, exist_ok=True)
-        with tarfile.open(path, "r:gz") as tar:
-            extracted_root = extracted.resolve()
-            for member in tar.getmembers():
-                destination = (extracted / member.name).resolve()
-                try:
-                    destination.relative_to(extracted_root)
-                except ValueError:
-                    raise ValueError(f"Unsafe archive member: {member.name}")
-            tar.extractall(extracted)
+        safe_extract_archive(path, extracted)
         candidates = [candidate for candidate in extracted.rglob("adapter_config.json")]
         if candidates:
             return str(candidates[0].parent)
@@ -271,6 +290,7 @@ def create_text_model_and_tokenizer(model_source: str):
     token = os.getenv("HF_TOKEN")
     tokenizer = AutoTokenizer.from_pretrained(
         model_source,
+        **model_revision_kwargs(model_source),
         trust_remote_code=trust_remote_code,
         token=token,
     )
@@ -280,6 +300,7 @@ def create_text_model_and_tokenizer(model_source: str):
     dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
     model = AutoModelForCausalLM.from_pretrained(
         model_source,
+        **model_revision_kwargs(model_source),
         trust_remote_code=trust_remote_code,
         token=token,
         torch_dtype=dtype if torch.cuda.is_available() else None,
@@ -294,6 +315,7 @@ def create_multimodal_model_and_processor(model_source: str):
     token = os.getenv("HF_TOKEN")
     processor = AutoProcessor.from_pretrained(
         model_source,
+        **model_revision_kwargs(model_source),
         trust_remote_code=trust_remote_code,
         token=token,
     )
@@ -303,6 +325,7 @@ def create_multimodal_model_and_processor(model_source: str):
     dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
     model = AutoModelForImageTextToText.from_pretrained(
         model_source,
+        **model_revision_kwargs(model_source),
         trust_remote_code=trust_remote_code,
         token=token,
         torch_dtype=dtype if torch.cuda.is_available() else None,
@@ -438,17 +461,19 @@ def main() -> None:
     OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
 
     rows = load_rows()
-    model_source = resolve_model_source()
-    model_loader = hp("model_loader", "causal_lm")
-    if model_loader == "image_text_to_text":
-        metrics = run_multimodal_training(rows, model_source)
-    else:
-        metrics = run_text_training(rows, model_source)
+    with TemporaryDirectory(prefix="tt-local-base-model-") as base_tmp:
+        model_source = resolve_model_source(Path(base_tmp))
+        model_loader = hp("model_loader", "causal_lm")
+        if model_loader == "image_text_to_text":
+            metrics = run_multimodal_training(rows, model_source)
+        else:
+            metrics = run_text_training(rows, model_source)
 
     archive_path = create_model_archive()
     output_metrics = {
         "training_rows": len(rows),
-        "model_source": model_source,
+        "model_source": str(BASE_MODEL_DIR) if BASE_MODEL_DIR.is_dir() and any(BASE_MODEL_DIR.iterdir()) else hp("base_model"),
+        "base_model_revision": hp("base_model_revision"),
         "model_loader": model_loader,
         "parent_model_artifact": hp("parent_model_artifact"),
         "train_runtime": round(time.time() - started, 3),