Track Kimi model metadata and cache path

packages/node/meshnet_node/torch_server.py

@@ -12,6 +12,7 @@ import urllib.error
 import urllib.parse
 import urllib.request
 import uuid
+from pathlib import Path
 from typing import Any
 
 from .model_backend import (
@@ -682,6 +683,7 @@ class TorchNodeServer:
         tracker_mode: bool | None = None,
         tracker_url: str | None = None,
         route_timeout: float = 30.0,
+        cache_dir: Path | None = None,
         debug: bool = False,
     ) -> None:
         self._host = host
@@ -691,11 +693,13 @@ class TorchNodeServer:
             shard_start,
             shard_end,
             quantization,
+            cache_dir,
         )
         # Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
         self._tracker_mode = tracker_mode if tracker_mode is not None else (shard_start == 0)
         self._tracker_url = tracker_url
         self._route_timeout = route_timeout
+        self._cache_dir = cache_dir
         self._debug = debug
         self._server: _TorchHTTPServer | None = None
         self._thread: threading.Thread | None = None
@@ -745,7 +749,10 @@ class TorchNodeServer:
             f"  [node] loading reassigned shard: {model_id} layers {shard_start}-{shard_end}",
             flush=True,
         )
-        new_backend = _load_backend(model_id, shard_start, shard_end, quantization)
+        try:
+            new_backend = _load_backend(model_id, shard_start, shard_end, quantization, self._cache_dir)
+        except TypeError:
+            new_backend = _load_backend(model_id, shard_start, shard_end, quantization)
         self._backend = new_backend
         self._tracker_mode = shard_start == 0
         if self._server is not None:
@@ -797,12 +804,13 @@ def _load_backend(
     shard_start: int,
     shard_end: int,
     quantization: str,
+    cache_dir: Path | None = None,
 ) -> TorchModelShard:
     from .model_backend import load_torch_shard
 
     quant = validate_quantization(quantization)
     try:
-        return load_torch_shard(model_id, shard_start, shard_end, quant)
+        return load_torch_shard(model_id, shard_start, shard_end, quant, cache_dir)
     except MissingModelDependencyError:
         raise
     except InsufficientVRAMError as exc: