Add embedded ModelHub auth fallback
This commit is contained in:
@@ -6,6 +6,7 @@
|
|||||||
**/.ipynb_checkpoints/
|
**/.ipynb_checkpoints/
|
||||||
|
|
||||||
# Never bake local credentials into the strategy image.
|
# Never bake local credentials into the strategy image.
|
||||||
|
Token
|
||||||
KEY.md
|
KEY.md
|
||||||
KEYS.md
|
KEYS.md
|
||||||
modelhub_submmit_api/KEY.md
|
modelhub_submmit_api/KEY.md
|
||||||
|
|||||||
1
.gitignore
vendored
1
.gitignore
vendored
@@ -3,6 +3,7 @@ __pycache__/
|
|||||||
.ipynb_checkpoints/
|
.ipynb_checkpoints/
|
||||||
|
|
||||||
# Local credentials.
|
# Local credentials.
|
||||||
|
Token
|
||||||
KEY.md
|
KEY.md
|
||||||
KEYS.md
|
KEYS.md
|
||||||
modelhub_submmit_api/KEY.md
|
modelhub_submmit_api/KEY.md
|
||||||
|
|||||||
@@ -15,11 +15,13 @@ submission poller in a child process.
|
|||||||
|
|
||||||
## Runtime Environment
|
## Runtime Environment
|
||||||
|
|
||||||
Set credentials with environment variables. Do not commit local key files.
|
The image includes a single-account submission token fallback for the agent
|
||||||
|
platform. Environment variables can override it without rebuilding the image.
|
||||||
|
|
||||||
- `MODELHUB_XC_TOKEN` or `XC_TOKEN`
|
- `MODELHUB_XC_TOKEN`, `XC_TOKEN`, or `MODELHUB_TOKEN` for ModelHub API authentication
|
||||||
|
- `MODELHUB_JWT_TOKEN` or `JWT_TOKEN` can be used instead when the platform provides a JWT
|
||||||
- `HF_TOKEN` optional, but recommended for Hugging Face API limits
|
- `HF_TOKEN` optional, but recommended for Hugging Face API limits
|
||||||
- `STRATEGY_ID` is expected to be injected by the ModelHub agent platform
|
- `STRATEGY_ID` is expected to be injected by the ModelHub agent platform and is attached to submissions for strategy attribution; it is not an API authentication token
|
||||||
|
|
||||||
Optional tuning:
|
Optional tuning:
|
||||||
|
|
||||||
|
|||||||
27
main.py
27
main.py
@@ -9,6 +9,8 @@ import time
|
|||||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
from modelhub_submmit_api.defaults import EMBEDDED_MODELHUB_XC_TOKEN
|
||||||
|
|
||||||
|
|
||||||
HOST = "0.0.0.0"
|
HOST = "0.0.0.0"
|
||||||
PORT = int(os.getenv("PORT", "8080"))
|
PORT = int(os.getenv("PORT", "8080"))
|
||||||
@@ -17,6 +19,18 @@ WORKER_SCRIPT = ROOT / "modelhub_submmit_api" / "poll_runner.py"
|
|||||||
|
|
||||||
shutdown_requested = False
|
shutdown_requested = False
|
||||||
worker: subprocess.Popen | None = None
|
worker: subprocess.Popen | None = None
|
||||||
|
config_error: str | None = None
|
||||||
|
|
||||||
|
|
||||||
|
def _has_modelhub_auth() -> bool:
|
||||||
|
return bool(
|
||||||
|
os.getenv("MODELHUB_XC_TOKEN")
|
||||||
|
or os.getenv("XC_TOKEN")
|
||||||
|
or os.getenv("MODELHUB_TOKEN")
|
||||||
|
or os.getenv("MODELHUB_JWT_TOKEN")
|
||||||
|
or os.getenv("JWT_TOKEN")
|
||||||
|
or EMBEDDED_MODELHUB_XC_TOKEN
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def _csv_args(env_name: str) -> list[str]:
|
def _csv_args(env_name: str) -> list[str]:
|
||||||
@@ -62,7 +76,8 @@ def _config() -> dict[str, object]:
|
|||||||
return {
|
return {
|
||||||
"strategy_id_present": bool(os.getenv("STRATEGY_ID")),
|
"strategy_id_present": bool(os.getenv("STRATEGY_ID")),
|
||||||
"hf_token_present": bool(os.getenv("HF_TOKEN")),
|
"hf_token_present": bool(os.getenv("HF_TOKEN")),
|
||||||
"modelhub_token_present": bool(os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN")),
|
"modelhub_auth_present": _has_modelhub_auth(),
|
||||||
|
"config_error": config_error,
|
||||||
"worker_running": worker is not None and worker.poll() is None,
|
"worker_running": worker is not None and worker.poll() is None,
|
||||||
"worker_return_code": None if worker is None else worker.poll(),
|
"worker_return_code": None if worker is None else worker.poll(),
|
||||||
}
|
}
|
||||||
@@ -71,6 +86,9 @@ def _config() -> dict[str, object]:
|
|||||||
class Handler(BaseHTTPRequestHandler):
|
class Handler(BaseHTTPRequestHandler):
|
||||||
def do_GET(self) -> None:
|
def do_GET(self) -> None:
|
||||||
if self.path == "/health":
|
if self.path == "/health":
|
||||||
|
if config_error:
|
||||||
|
self._send_json({"status": "config_error", "config": _config()}, status=500)
|
||||||
|
return
|
||||||
if worker is not None and worker.poll() is not None:
|
if worker is not None and worker.poll() is not None:
|
||||||
self._send_json({"status": "worker_exited", "config": _config()}, status=500)
|
self._send_json({"status": "worker_exited", "config": _config()}, status=500)
|
||||||
return
|
return
|
||||||
@@ -116,12 +134,15 @@ def _handle_signal(signum: int, _frame: object) -> None:
|
|||||||
|
|
||||||
|
|
||||||
def main() -> int:
|
def main() -> int:
|
||||||
global worker
|
global config_error, worker
|
||||||
signal.signal(signal.SIGTERM, _handle_signal)
|
signal.signal(signal.SIGTERM, _handle_signal)
|
||||||
signal.signal(signal.SIGINT, _handle_signal)
|
signal.signal(signal.SIGINT, _handle_signal)
|
||||||
|
|
||||||
start_worker = os.getenv("MODELHUB_AGENT_START_WORKER", "1").strip().lower() not in {"0", "false", "no"}
|
start_worker = os.getenv("MODELHUB_AGENT_START_WORKER", "1").strip().lower() not in {"0", "false", "no"}
|
||||||
if start_worker:
|
if start_worker and not _has_modelhub_auth():
|
||||||
|
config_error = "missing ModelHub auth: set MODELHUB_XC_TOKEN/XC_TOKEN or MODELHUB_JWT_TOKEN/JWT_TOKEN"
|
||||||
|
print(config_error, flush=True)
|
||||||
|
elif start_worker:
|
||||||
cmd = _worker_command()
|
cmd = _worker_command()
|
||||||
print("starting submission worker: " + " ".join(cmd), flush=True)
|
print("starting submission worker: " + " ".join(cmd), flush=True)
|
||||||
worker = subprocess.Popen(cmd, cwd=str(ROOT))
|
worker = subprocess.Popen(cmd, cwd=str(ROOT))
|
||||||
|
|||||||
1
modelhub_submmit_api/defaults.py
Normal file
1
modelhub_submmit_api/defaults.py
Normal file
@@ -0,0 +1 @@
|
|||||||
|
EMBEDDED_MODELHUB_XC_TOKEN = "a14776f6e7ad4c04a1710260613c294c"
|
||||||
@@ -5,6 +5,7 @@ from datetime import datetime, timedelta
|
|||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from common import format_modelhub_datetime, parse_datetime
|
from common import format_modelhub_datetime, parse_datetime
|
||||||
|
from defaults import EMBEDDED_MODELHUB_XC_TOKEN
|
||||||
from http_json import HttpJsonError, JsonHttpClient
|
from http_json import HttpJsonError, JsonHttpClient
|
||||||
|
|
||||||
|
|
||||||
@@ -25,10 +26,21 @@ class ModelHubClient:
|
|||||||
retries: int = 2,
|
retries: int = 2,
|
||||||
http_client: JsonHttpClient | None = None,
|
http_client: JsonHttpClient | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.token = token or os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN")
|
self.token = (
|
||||||
if not self.token and not os.getenv("STRATEGY_ID") and http_client is None:
|
token
|
||||||
raise ValueError("ModelHub token is required. Set MODELHUB_XC_TOKEN or XC_TOKEN.")
|
or os.getenv("MODELHUB_XC_TOKEN")
|
||||||
default_headers = {"Xc-Token": self.token} if self.token else {}
|
or os.getenv("XC_TOKEN")
|
||||||
|
or os.getenv("MODELHUB_TOKEN")
|
||||||
|
or EMBEDDED_MODELHUB_XC_TOKEN
|
||||||
|
)
|
||||||
|
self.jwt_token = os.getenv("MODELHUB_JWT_TOKEN") or os.getenv("JWT_TOKEN")
|
||||||
|
if not self.token and not self.jwt_token and http_client is None:
|
||||||
|
raise ValueError("ModelHub token is required. Set MODELHUB_XC_TOKEN/XC_TOKEN or MODELHUB_JWT_TOKEN/JWT_TOKEN.")
|
||||||
|
default_headers = {}
|
||||||
|
if self.token:
|
||||||
|
default_headers["Xc-Token"] = self.token
|
||||||
|
if self.jwt_token:
|
||||||
|
default_headers["Authorization"] = f"Bearer {self.jwt_token}"
|
||||||
self.http_client = http_client or JsonHttpClient(
|
self.http_client = http_client or JsonHttpClient(
|
||||||
base_url=base_url,
|
base_url=base_url,
|
||||||
default_headers=default_headers,
|
default_headers=default_headers,
|
||||||
@@ -227,4 +239,3 @@ def is_active_task(task: dict[str, Any]) -> bool:
|
|||||||
if status in TERMINAL_TASK_STATUSES:
|
if status in TERMINAL_TASK_STATUSES:
|
||||||
return False
|
return False
|
||||||
return True
|
return True
|
||||||
|
|
||||||
|
|||||||
@@ -4,8 +4,12 @@ import argparse
|
|||||||
import os
|
import os
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
from defaults import EMBEDDED_MODELHUB_XC_TOKEN
|
||||||
|
|
||||||
|
|
||||||
DEFAULT_KEY_PATH = Path("KEY.md")
|
DEFAULT_KEY_PATH = Path("KEY.md")
|
||||||
|
XC_TOKEN_ENV_NAMES = ("MODELHUB_XC_TOKEN", "XC_TOKEN", "MODELHUB_TOKEN")
|
||||||
|
JWT_TOKEN_ENV_NAMES = ("MODELHUB_JWT_TOKEN", "JWT_TOKEN")
|
||||||
|
|
||||||
|
|
||||||
def load_key_file(path: Path) -> dict[str, str]:
|
def load_key_file(path: Path) -> dict[str, str]:
|
||||||
@@ -21,6 +25,14 @@ def load_key_file(path: Path) -> dict[str, str]:
|
|||||||
return loaded
|
return loaded
|
||||||
|
|
||||||
|
|
||||||
|
def first_value(values: dict[str, str], names: tuple[str, ...]) -> str | None:
|
||||||
|
for name in names:
|
||||||
|
value = os.getenv(name) or values.get(name)
|
||||||
|
if value:
|
||||||
|
return value.strip()
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
def ensure_tokens(args: argparse.Namespace) -> None:
|
def ensure_tokens(args: argparse.Namespace) -> None:
|
||||||
primary_key_path = Path(getattr(args, "key_path", DEFAULT_KEY_PATH))
|
primary_key_path = Path(getattr(args, "key_path", DEFAULT_KEY_PATH))
|
||||||
if not primary_key_path.exists():
|
if not primary_key_path.exists():
|
||||||
@@ -34,7 +46,10 @@ def ensure_tokens(args: argparse.Namespace) -> None:
|
|||||||
|
|
||||||
modelhub_token = getattr(args, "modelhub_token", None)
|
modelhub_token = getattr(args, "modelhub_token", None)
|
||||||
if not modelhub_token:
|
if not modelhub_token:
|
||||||
modelhub_token = os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN") or values.get("MODELHUB_XC_TOKEN") or values.get("XC_TOKEN")
|
modelhub_token = first_value(values, XC_TOKEN_ENV_NAMES)
|
||||||
|
if not modelhub_token:
|
||||||
|
modelhub_token = EMBEDDED_MODELHUB_XC_TOKEN
|
||||||
|
jwt_token = first_value(values, JWT_TOKEN_ENV_NAMES)
|
||||||
|
|
||||||
args.modelhub_token = modelhub_token
|
args.modelhub_token = modelhub_token
|
||||||
|
|
||||||
@@ -43,5 +58,7 @@ def ensure_tokens(args: argparse.Namespace) -> None:
|
|||||||
if args.modelhub_token:
|
if args.modelhub_token:
|
||||||
os.environ["MODELHUB_XC_TOKEN"] = args.modelhub_token
|
os.environ["MODELHUB_XC_TOKEN"] = args.modelhub_token
|
||||||
os.environ["XC_TOKEN"] = args.modelhub_token
|
os.environ["XC_TOKEN"] = args.modelhub_token
|
||||||
if not args.modelhub_token and not os.getenv("STRATEGY_ID"):
|
if jwt_token:
|
||||||
raise ValueError("XC_TOKEN is required, either via environment or KEY.md")
|
os.environ["MODELHUB_JWT_TOKEN"] = jwt_token
|
||||||
|
if not args.modelhub_token and not jwt_token:
|
||||||
|
raise ValueError("MODELHUB_XC_TOKEN/XC_TOKEN or MODELHUB_JWT_TOKEN/JWT_TOKEN is required")
|
||||||
|
|||||||
Reference in New Issue
Block a user