update load

src/twinkle/server/tinker/common/megatron_model.py

@@ -175,13 +175,13 @@ def load(self, checkpoint_dir: str, **kwargs):
 
         # Create checkpoint manager with the token
         checkpoint_manager = create_checkpoint_manager(token)
-        
+
         # Use resolve_load_path to handle path resolution
         resolved = checkpoint_manager.resolve_load_path(checkpoint_dir)
 
         if resolved.is_twinkle_path:
             # Load from twinkle checkpoint
-            return super().load(name=resolved.checkpoint_name, output_dir=str(resolved.checkpoint_dir), **kwargs)
+            return super().load(name=resolved.checkpoint_name, output_dir=resolved.checkpoint_dir, **kwargs)
         else:
             # Load from hub
             return super().load(name=resolved.checkpoint_name, **kwargs)

src/twinkle/server/tinker/common/transformers_model.py

@@ -137,7 +137,7 @@ def load(self, checkpoint_dir: str, **kwargs):
 
         if resolved.is_twinkle_path:
             # Load from twinkle checkpoint
-            return super().load(name=resolved.checkpoint_name, output_dir=str(resolved.checkpoint_dir), **kwargs)
+            return super().load(name=resolved.checkpoint_name, output_dir=resolved.checkpoint_dir, **kwargs)
         else:
             # Load from hub
             return super().load(name=resolved.checkpoint_name, **kwargs)

src/twinkle/server/utils/io_utils.py

@@ -912,18 +912,16 @@ def resolve_load_path(self, path: str, validate_exists: bool = True) -> Resolved
                         f"Checkpoint not found or access denied: {path}"
                     )
 
-            # Get the checkpoint directory
-            checkpoint_dir = str(self.get_ckpt_dir(training_run_id, checkpoint_id))
+            # Get the checkpoint directory parent path (no checkpoint name in the path)
+            checkpoint_dir = self.get_ckpt_dir(training_run_id, checkpoint_id).parent
 
             if validate_exists:
-                # Verify the directory exists
-                from pathlib import Path as PathLib
-                if not PathLib(checkpoint_dir).exists():
+                if not checkpoint_dir.exists():
                     raise ValueError(f"Checkpoint directory not found: {checkpoint_dir}")
 
             return ResolvedLoadPath(
                 checkpoint_name=checkpoint_name,
-                checkpoint_dir=checkpoint_dir,
+                checkpoint_dir=checkpoint_dir.as_posix(),
                 is_twinkle_path=True,
                 training_run_id=training_run_id,
                 checkpoint_id=checkpoint_id

src/twinkle/server/utils/validation.py

@@ -22,7 +22,7 @@ async def verify_request_token(request: Request, call_next):
     Returns:
         JSONResponse with error if validation fails, otherwise the response from call_next
     """
-    authorization = request.headers.get("Authorization")
+    authorization = request.headers.get("Twinkle-Authorization")
     token = authorization[7:] if authorization and authorization.startswith("Bearer ") else authorization
     if not is_token_valid(token):
         return JSONResponse(status_code=403, content={"detail": "Invalid token"})

src/twinkle_client/__init__.py

@@ -29,6 +29,7 @@ def init_tinker_compat_client(base_url: Optional[str] = None, api_key: Optional[
     default_headers = {
         "X-Ray-Serve-Request-Id": get_request_id(),
         "Authorization": 'Bearer ' + api_key,
+        "Twinkle-Authorization": 'Bearer ' + api_key,  # For server compatibility
     } | kwargs.pop("default_headers", {})
 
     service_client = ServiceClient(base_url=base_url, api_key=api_key, default_headers=default_headers, **kwargs)

src/twinkle_client/http/http_utils.py

@@ -17,6 +17,7 @@ def _build_headers(additional_headers: Optional[Dict[str, str]] = None) -> Dict[
     headers = {
         "X-Ray-Serve-Request-Id": get_request_id(),
         "Authorization": 'Bearer ' + get_api_key(),
+        "Twinkle-Authorization": 'Bearer ' + get_api_key(),  # For server compatibility
     }
 
     if additional_headers: