fix: coordinate concurrent account capacity filling
This commit is contained in:
@@ -27,3 +27,5 @@ modelhub_submmit_api/poll_runs/
|
|||||||
modelhub_submmit_api/ledger/
|
modelhub_submmit_api/ledger/
|
||||||
modelhub_submmit_api/outcomes/
|
modelhub_submmit_api/outcomes/
|
||||||
modelhub_submmit_api/history/
|
modelhub_submmit_api/history/
|
||||||
|
.modelhub_state/
|
||||||
|
modelhub_submmit_api/.modelhub_state/
|
||||||
|
|||||||
2
.gitignore
vendored
2
.gitignore
vendored
@@ -24,3 +24,5 @@ modelhub_submmit_api/poll_runs/
|
|||||||
modelhub_submmit_api/ledger/
|
modelhub_submmit_api/ledger/
|
||||||
modelhub_submmit_api/outcomes/
|
modelhub_submmit_api/outcomes/
|
||||||
modelhub_submmit_api/history/
|
modelhub_submmit_api/history/
|
||||||
|
.modelhub_state/
|
||||||
|
modelhub_submmit_api/.modelhub_state/
|
||||||
|
|||||||
26
README.md
26
README.md
@@ -26,15 +26,33 @@ the image.
|
|||||||
|
|
||||||
Optional tuning:
|
Optional tuning:
|
||||||
|
|
||||||
- `MODELHUB_AGENT_POLL_INTERVAL_SECONDS` default `300`
|
- `MODELHUB_AGENT_POLL_INTERVAL_SECONDS` default `15`
|
||||||
- `MODELHUB_AGENT_IDLE_INTERVAL_SECONDS` default `600`
|
- `MODELHUB_AGENT_IDLE_INTERVAL_SECONDS` default `60`
|
||||||
- `MODELHUB_AGENT_POST_CYCLE_COOLDOWN_SECONDS` default `30`
|
- `MODELHUB_AGENT_POST_CYCLE_COOLDOWN_SECONDS` default `2`
|
||||||
- `MODELHUB_AGENT_MAX_SUBMITS_PER_RUN` default `5`
|
- `MODELHUB_AGENT_MAX_SUBMITS_PER_RUN` default `0` (fill all currently available slots)
|
||||||
|
- `MODELHUB_AGENT_ACTIVE_TASK_CAP` default `100` per account
|
||||||
|
- `MODELHUB_AGENT_ACTIVE_COUNTS_TTL_SECONDS` default `15`
|
||||||
|
- `MODELHUB_AGENT_RESERVATION_TTL_SECONDS` default `120`
|
||||||
|
- `MODELHUB_AGENT_INSTANCE_ID` optional stable worker identity used to spread concurrent agents across accounts and candidates
|
||||||
|
- `MODELHUB_AGENT_CLAIMS_PATH` default `.modelhub_state/submission_claims.jsonl`
|
||||||
- `MODELHUB_AGENT_DAILY_TARGET`
|
- `MODELHUB_AGENT_DAILY_TARGET`
|
||||||
- `MODELHUB_AGENT_MIN_DOWNLOADS`
|
- `MODELHUB_AGENT_MIN_DOWNLOADS`
|
||||||
- `MODELHUB_AGENT_GPUS`
|
- `MODELHUB_AGENT_GPUS`
|
||||||
- `MODELHUB_AGENT_EXTRA_ARGS`
|
- `MODELHUB_AGENT_EXTRA_ARGS`
|
||||||
|
|
||||||
|
## Concurrent Agents
|
||||||
|
|
||||||
|
The token pool keeps a local reservation for every in-flight submission, so a
|
||||||
|
lagging platform count cannot send all concurrent requests to the same account.
|
||||||
|
If another process fills an account first, the submission is retried immediately
|
||||||
|
against another account with capacity.
|
||||||
|
|
||||||
|
Workers that share a filesystem also coordinate model/GPU claims through
|
||||||
|
`.modelhub_state/submission_claims.jsonl`. Workers in isolated containers use
|
||||||
|
different candidate ordering (derived from `STRATEGY_ID`, instance ID, or
|
||||||
|
hostname), which reduces duplicate work while the platform remains the final
|
||||||
|
authority for account capacity and model/GPU uniqueness.
|
||||||
|
|
||||||
## Deploy
|
## Deploy
|
||||||
|
|
||||||
Create a tag and submit the repository URL plus tag in "我的适配智能体".
|
Create a tag and submit the repository URL plus tag in "我的适配智能体".
|
||||||
|
|||||||
8
main.py
8
main.py
@@ -46,13 +46,13 @@ def _worker_command() -> list[str]:
|
|||||||
"-u",
|
"-u",
|
||||||
str(WORKER_SCRIPT),
|
str(WORKER_SCRIPT),
|
||||||
"--poll-interval-seconds",
|
"--poll-interval-seconds",
|
||||||
os.getenv("MODELHUB_AGENT_POLL_INTERVAL_SECONDS", "300"),
|
os.getenv("MODELHUB_AGENT_POLL_INTERVAL_SECONDS", "15"),
|
||||||
"--idle-interval-seconds",
|
"--idle-interval-seconds",
|
||||||
os.getenv("MODELHUB_AGENT_IDLE_INTERVAL_SECONDS", "600"),
|
os.getenv("MODELHUB_AGENT_IDLE_INTERVAL_SECONDS", "60"),
|
||||||
"--post-cycle-cooldown-seconds",
|
"--post-cycle-cooldown-seconds",
|
||||||
os.getenv("MODELHUB_AGENT_POST_CYCLE_COOLDOWN_SECONDS", "30"),
|
os.getenv("MODELHUB_AGENT_POST_CYCLE_COOLDOWN_SECONDS", "2"),
|
||||||
"--max-submits-per-run",
|
"--max-submits-per-run",
|
||||||
os.getenv("MODELHUB_AGENT_MAX_SUBMITS_PER_RUN", "5"),
|
os.getenv("MODELHUB_AGENT_MAX_SUBMITS_PER_RUN", "0"),
|
||||||
"--skip-history-archive",
|
"--skip-history-archive",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|||||||
@@ -74,6 +74,9 @@ bash run_poll.sh --dry-run
|
|||||||
- The runner auto-discovers all safe GPU/template combinations from the public submit catalog.
|
- The runner auto-discovers all safe GPU/template combinations from the public submit catalog.
|
||||||
- Each model can be submitted at most once per GPU.
|
- Each model can be submitted at most once per GPU.
|
||||||
- Multiple ModelHub tokens are pooled and used to route submissions to the account with available async capacity.
|
- Multiple ModelHub tokens are pooled and used to route submissions to the account with available async capacity.
|
||||||
|
- Concurrent submissions reserve account slots locally, and an account-capacity race automatically falls through to another account.
|
||||||
|
- Concurrent local processes claim model/GPU pairs in `.modelhub_state/submission_claims.jsonl`; shared ledger, history, and outcome files use process locks and atomic replacement.
|
||||||
|
- Isolated agent containers diversify candidate order by instance identity to reduce cross-container duplicate submissions.
|
||||||
- History stats are online-only. The local ledger is used for local accounting, but platform history is only used after the local ledger reaches the configured threshold.
|
- History stats are online-only. The local ledger is used for local accounting, but platform history is only used after the local ledger reaches the configured threshold.
|
||||||
- The default history threshold is `500` records.
|
- The default history threshold is `500` records.
|
||||||
|
|
||||||
@@ -96,15 +99,15 @@ Common flags:
|
|||||||
|
|
||||||
`run_poll.sh` adds:
|
`run_poll.sh` adds:
|
||||||
|
|
||||||
- `--poll-interval-seconds`: sleep when all accounts are saturated
|
- `--poll-interval-seconds`: sleep when all accounts are saturated (default 15)
|
||||||
- `--idle-interval-seconds`: sleep when a cycle submits nothing
|
- `--idle-interval-seconds`: sleep when a cycle submits nothing (default 60)
|
||||||
- `--max-scan-models`: hard cap on scanned HF models for this cycle (0 = auto)
|
- `--max-scan-models`: hard cap on scanned HF models for this cycle (0 = auto)
|
||||||
- `--scan-multiplier`: multiplier used for auto scan cap derivation from quota/queue capacity
|
- `--scan-multiplier`: multiplier used for auto scan cap derivation from quota/queue capacity
|
||||||
- `--max-submits-per-run`: max tasks to submit per poll cycle (0 = unlimited)
|
- `--max-submits-per-run`: max tasks to submit per poll cycle (0 = unlimited)
|
||||||
- `--skip-outcome-sync`: skip outcome sync before scanning
|
- `--skip-outcome-sync`: skip outcome sync before scanning
|
||||||
- `--skip-history-archive`: skip history archive download for this cycle
|
- `--skip-history-archive`: skip history archive download for this cycle
|
||||||
- `--submit-concurrency`: concurrent task submission calls used by each cycle (0 = auto)
|
- `--submit-concurrency`: concurrent task submission calls used by each cycle (0 = auto)
|
||||||
- `--post-cycle-cooldown-seconds`: pause after a successful cycle before next cycle (default 0)
|
- `--post-cycle-cooldown-seconds`: pause after a successful cycle before next cycle (default 2)
|
||||||
- `--max-cycles`: optional hard stop for testing or batch windows
|
- `--max-cycles`: optional hard stop for testing or batch windows
|
||||||
|
|
||||||
## Output
|
## Output
|
||||||
@@ -140,3 +143,6 @@ python3 -m unittest discover -s tests -v
|
|||||||
- The platform still enforces per-account async capacity limits, so the poller can
|
- The platform still enforces per-account async capacity limits, so the poller can
|
||||||
keep the queue close to full but cannot override the platform cap.
|
keep the queue close to full but cannot override the platform cap.
|
||||||
- `bash run_poll.sh` now defaults to unlimited mode and keeps refilling until you stop the process manually.
|
- `bash run_poll.sh` now defaults to unlimited mode and keeps refilling until you stop the process manually.
|
||||||
|
- Queue polling defaults to 15 seconds, successful-cycle cooldown to 2 seconds,
|
||||||
|
and per-cycle submissions to all available slots. Override these values when
|
||||||
|
the platform requires a lower request rate.
|
||||||
|
|||||||
@@ -1,9 +1,15 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import fcntl
|
||||||
|
import hashlib
|
||||||
import json
|
import json
|
||||||
|
import os
|
||||||
|
import socket
|
||||||
|
import tempfile
|
||||||
|
from contextlib import contextmanager
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any, Callable, Iterator
|
||||||
|
|
||||||
|
|
||||||
def ensure_utc(value: datetime) -> datetime:
|
def ensure_utc(value: datetime) -> datetime:
|
||||||
@@ -54,16 +60,57 @@ def json_dumps(value: Any) -> str:
|
|||||||
return json.dumps(value, ensure_ascii=False, sort_keys=True)
|
return json.dumps(value, ensure_ascii=False, sort_keys=True)
|
||||||
|
|
||||||
|
|
||||||
def read_json(path: Path) -> Any:
|
def runtime_instance_id() -> str:
|
||||||
return json.loads(path.read_text(encoding="utf-8"))
|
"""Return a non-secret identity that differs between concurrent workers."""
|
||||||
|
base = (
|
||||||
|
os.getenv("MODELHUB_AGENT_INSTANCE_ID")
|
||||||
|
or os.getenv("STRATEGY_ID")
|
||||||
|
or os.getenv("HOSTNAME")
|
||||||
|
or socket.gethostname()
|
||||||
|
)
|
||||||
|
raw = f"{base}:{os.getpid()}"
|
||||||
|
return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:20]
|
||||||
|
|
||||||
|
|
||||||
def write_json(path: Path, value: Any) -> None:
|
def _lock_path(path: Path) -> Path:
|
||||||
|
return path.with_name(f".{path.name}.lock")
|
||||||
|
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def file_lock(path: Path, *, shared: bool = False) -> Iterator[None]:
|
||||||
|
"""Coordinate access to a path between processes on the same filesystem."""
|
||||||
path.parent.mkdir(parents=True, exist_ok=True)
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
path.write_text(json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
lock_path = _lock_path(path)
|
||||||
|
with lock_path.open("a+b") as handle:
|
||||||
|
fcntl.flock(handle.fileno(), fcntl.LOCK_SH if shared else fcntl.LOCK_EX)
|
||||||
|
try:
|
||||||
|
yield
|
||||||
|
finally:
|
||||||
|
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
|
||||||
|
|
||||||
|
|
||||||
def read_jsonl(path: Path) -> list[dict[str, Any]]:
|
def _atomic_write_text(path: Path, content: str) -> None:
|
||||||
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
descriptor, temporary_name = tempfile.mkstemp(
|
||||||
|
dir=str(path.parent),
|
||||||
|
prefix=f".{path.name}.",
|
||||||
|
suffix=".tmp",
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
|
||||||
|
handle.write(content)
|
||||||
|
handle.flush()
|
||||||
|
os.fsync(handle.fileno())
|
||||||
|
os.replace(temporary_name, path)
|
||||||
|
except Exception:
|
||||||
|
try:
|
||||||
|
os.unlink(temporary_name)
|
||||||
|
except FileNotFoundError:
|
||||||
|
pass
|
||||||
|
raise
|
||||||
|
|
||||||
|
|
||||||
|
def _read_jsonl_unlocked(path: Path) -> list[dict[str, Any]]:
|
||||||
if not path.exists():
|
if not path.exists():
|
||||||
return []
|
return []
|
||||||
rows: list[dict[str, Any]] = []
|
rows: list[dict[str, Any]] = []
|
||||||
@@ -74,13 +121,45 @@ def read_jsonl(path: Path) -> list[dict[str, Any]]:
|
|||||||
return rows
|
return rows
|
||||||
|
|
||||||
|
|
||||||
|
def _jsonl_content(rows: list[dict[str, Any]]) -> str:
|
||||||
|
return "".join(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n" for row in rows)
|
||||||
|
|
||||||
|
|
||||||
|
def read_json(path: Path) -> Any:
|
||||||
|
with file_lock(path, shared=True):
|
||||||
|
return json.loads(path.read_text(encoding="utf-8"))
|
||||||
|
|
||||||
|
|
||||||
|
def write_json(path: Path, value: Any) -> None:
|
||||||
|
content = json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
|
||||||
|
with file_lock(path):
|
||||||
|
_atomic_write_text(path, content)
|
||||||
|
|
||||||
|
|
||||||
|
def read_jsonl(path: Path) -> list[dict[str, Any]]:
|
||||||
|
with file_lock(path, shared=True):
|
||||||
|
return _read_jsonl_unlocked(path)
|
||||||
|
|
||||||
|
|
||||||
def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
|
def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
|
||||||
path.parent.mkdir(parents=True, exist_ok=True)
|
with file_lock(path):
|
||||||
content = "".join(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n" for row in rows)
|
_atomic_write_text(path, _jsonl_content(rows))
|
||||||
path.write_text(content, encoding="utf-8")
|
|
||||||
|
|
||||||
|
|
||||||
def append_jsonl(path: Path, row: dict[str, Any]) -> None:
|
def append_jsonl(path: Path, row: dict[str, Any]) -> None:
|
||||||
path.parent.mkdir(parents=True, exist_ok=True)
|
line = json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n"
|
||||||
with path.open("a", encoding="utf-8") as handle:
|
with file_lock(path):
|
||||||
handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")
|
with path.open("a", encoding="utf-8") as handle:
|
||||||
|
handle.write(line)
|
||||||
|
handle.flush()
|
||||||
|
|
||||||
|
|
||||||
|
def update_jsonl(
|
||||||
|
path: Path,
|
||||||
|
updater: Callable[[list[dict[str, Any]]], list[dict[str, Any]]],
|
||||||
|
) -> list[dict[str, Any]]:
|
||||||
|
"""Atomically read, update, and replace a JSONL file under one lock."""
|
||||||
|
with file_lock(path):
|
||||||
|
updated = updater(_read_jsonl_unlocked(path))
|
||||||
|
_atomic_write_text(path, _jsonl_content(updated))
|
||||||
|
return updated
|
||||||
|
|||||||
@@ -8,10 +8,11 @@ from typing import Any, Callable
|
|||||||
|
|
||||||
from common import utc_now, write_json
|
from common import utc_now, write_json
|
||||||
from hf_discovery import HuggingFaceDiscovery
|
from hf_discovery import HuggingFaceDiscovery
|
||||||
from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR, run_submission
|
from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR, make_run_dir, run_submission
|
||||||
from modelhub_client import ModelHubClient, ModelHubClientPool
|
from modelhub_client import ModelHubClient, ModelHubClientPool
|
||||||
from outcome_tracker import OutcomeTracker
|
from outcome_tracker import OutcomeTracker
|
||||||
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
|
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
|
||||||
|
from submission_claims import DEFAULT_CLAIMS_PATH
|
||||||
from template_selector import TemplateSelector
|
from template_selector import TemplateSelector
|
||||||
|
|
||||||
|
|
||||||
@@ -81,6 +82,11 @@ def build_parser() -> argparse.ArgumentParser:
|
|||||||
parser.add_argument("--key-path", default=str(DEFAULT_KEY_PATH), help="Path to KEY.md containing MODELSCOPE_TOKEN/XC_TOKEN")
|
parser.add_argument("--key-path", default=str(DEFAULT_KEY_PATH), help="Path to KEY.md containing MODELSCOPE_TOKEN/XC_TOKEN")
|
||||||
parser.add_argument("--runs-dir", default=str(DEFAULT_RUNS_DIR), help=argparse.SUPPRESS)
|
parser.add_argument("--runs-dir", default=str(DEFAULT_RUNS_DIR), help=argparse.SUPPRESS)
|
||||||
parser.add_argument("--ledger-path", default=str(DEFAULT_LEDGER_PATH), help=argparse.SUPPRESS)
|
parser.add_argument("--ledger-path", default=str(DEFAULT_LEDGER_PATH), help=argparse.SUPPRESS)
|
||||||
|
parser.add_argument(
|
||||||
|
"--claims-path",
|
||||||
|
default=os.getenv("MODELHUB_AGENT_CLAIMS_PATH", str(DEFAULT_CLAIMS_PATH)),
|
||||||
|
help=argparse.SUPPRESS,
|
||||||
|
)
|
||||||
parser.add_argument("--history-archive-path", default="history/platform_tasks.jsonl", help=argparse.SUPPRESS)
|
parser.add_argument("--history-archive-path", default="history/platform_tasks.jsonl", help=argparse.SUPPRESS)
|
||||||
parser.add_argument("--history-archive-limit", type=int, default=5000, help=argparse.SUPPRESS)
|
parser.add_argument("--history-archive-limit", type=int, default=5000, help=argparse.SUPPRESS)
|
||||||
parser.add_argument("--daily-runs-dir", default=str(DEFAULT_DAILY_RUNS_DIR), help=argparse.SUPPRESS)
|
parser.add_argument("--daily-runs-dir", default=str(DEFAULT_DAILY_RUNS_DIR), help=argparse.SUPPRESS)
|
||||||
@@ -117,6 +123,7 @@ def make_wave_namespace(base_args: argparse.Namespace, wave: WaveSpec) -> argpar
|
|||||||
runs_dir=base_args.runs_dir,
|
runs_dir=base_args.runs_dir,
|
||||||
ledger_path=base_args.ledger_path,
|
ledger_path=base_args.ledger_path,
|
||||||
outcomes_path=getattr(base_args, "outcomes_path", "outcomes/submissions.jsonl"),
|
outcomes_path=getattr(base_args, "outcomes_path", "outcomes/submissions.jsonl"),
|
||||||
|
claims_path=getattr(base_args, "claims_path", str(DEFAULT_CLAIMS_PATH)),
|
||||||
history_archive_path=base_args.history_archive_path,
|
history_archive_path=base_args.history_archive_path,
|
||||||
history_archive_limit=base_args.history_archive_limit,
|
history_archive_limit=base_args.history_archive_limit,
|
||||||
hf_base_url=base_args.hf_base_url,
|
hf_base_url=base_args.hf_base_url,
|
||||||
@@ -132,7 +139,7 @@ def run_daily_batches(
|
|||||||
now=None,
|
now=None,
|
||||||
run_fn: Callable[..., dict[str, Any]] = run_submission,
|
run_fn: Callable[..., dict[str, Any]] = run_submission,
|
||||||
hf_discovery: HuggingFaceDiscovery | None = None,
|
hf_discovery: HuggingFaceDiscovery | None = None,
|
||||||
modelhub_client: ModelHubClient | None = None,
|
modelhub_client: ModelHubClient | ModelHubClientPool | None = None,
|
||||||
template_selector: TemplateSelector | None = None,
|
template_selector: TemplateSelector | None = None,
|
||||||
outcome_tracker: OutcomeTracker | None = None,
|
outcome_tracker: OutcomeTracker | None = None,
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
@@ -140,15 +147,12 @@ def run_daily_batches(
|
|||||||
hf_discovery = hf_discovery or HuggingFaceDiscovery(base_url=base_args.hf_base_url)
|
hf_discovery = hf_discovery or HuggingFaceDiscovery(base_url=base_args.hf_base_url)
|
||||||
if modelhub_client is None:
|
if modelhub_client is None:
|
||||||
modelhub_tokens = list(getattr(base_args, "modelhub_tokens", None) or ([] if not base_args.modelhub_token else [base_args.modelhub_token]))
|
modelhub_tokens = list(getattr(base_args, "modelhub_tokens", None) or ([] if not base_args.modelhub_token else [base_args.modelhub_token]))
|
||||||
if len(modelhub_tokens) > 1:
|
token_values: list[str | None] = modelhub_tokens or [base_args.modelhub_token]
|
||||||
clients = [ModelHubClient(token=token, base_url=base_args.modelhub_base_url) for token in modelhub_tokens]
|
clients = [ModelHubClient(token=token, base_url=base_args.modelhub_base_url) for token in token_values]
|
||||||
modelhub_client = ModelHubClientPool(clients)
|
modelhub_client = ModelHubClientPool(clients)
|
||||||
else:
|
|
||||||
modelhub_client = ModelHubClient(token=base_args.modelhub_token, base_url=base_args.modelhub_base_url)
|
|
||||||
template_selector = template_selector or TemplateSelector()
|
template_selector = template_selector or TemplateSelector()
|
||||||
|
|
||||||
daily_run_dir = Path(base_args.daily_runs_dir) / now.strftime("%Y%m%dT%H%M%SZ")
|
daily_run_dir = make_run_dir(Path(base_args.daily_runs_dir), now)
|
||||||
daily_run_dir.mkdir(parents=True, exist_ok=True)
|
|
||||||
log(f"[daily] daily_run_dir={daily_run_dir}")
|
log(f"[daily] daily_run_dir={daily_run_dir}")
|
||||||
log(
|
log(
|
||||||
f"[daily] target={base_args.daily_target} rounds={base_args.rounds} "
|
f"[daily] target={base_args.daily_target} rounds={base_args.rounds} "
|
||||||
|
|||||||
@@ -5,7 +5,7 @@ from datetime import datetime
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from common import append_jsonl, ensure_utc, parse_datetime, read_jsonl, write_jsonl
|
from common import append_jsonl, ensure_utc, parse_datetime, read_jsonl, update_jsonl
|
||||||
|
|
||||||
|
|
||||||
WAITING_STATUSES = {"waiting", "queued"}
|
WAITING_STATUSES = {"waiting", "queued"}
|
||||||
@@ -32,10 +32,10 @@ def update_history_archive(
|
|||||||
*,
|
*,
|
||||||
limit: int = HISTORY_ARCHIVE_LIMIT_DEFAULT,
|
limit: int = HISTORY_ARCHIVE_LIMIT_DEFAULT,
|
||||||
) -> list[dict[str, Any]]:
|
) -> list[dict[str, Any]]:
|
||||||
merged = merge_history_records(load_history_archive(path), tasks, limit=limit)
|
return update_jsonl(
|
||||||
path.parent.mkdir(parents=True, exist_ok=True)
|
path,
|
||||||
write_jsonl(path, merged)
|
lambda existing: merge_history_records(existing, tasks, limit=limit),
|
||||||
return merged
|
)
|
||||||
|
|
||||||
|
|
||||||
def merge_history_records(
|
def merge_history_records(
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ from datetime import timedelta
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from common import parse_datetime, utc_now, write_json, write_jsonl
|
from common import parse_datetime, runtime_instance_id, utc_now, write_json, write_jsonl
|
||||||
from hf_discovery import HuggingFaceDiscovery
|
from hf_discovery import HuggingFaceDiscovery
|
||||||
from history_stats import (
|
from history_stats import (
|
||||||
append_ledger_entry,
|
append_ledger_entry,
|
||||||
@@ -19,6 +19,7 @@ from history_stats import (
|
|||||||
from modelhub_client import ModelHubAPIError, ModelHubClient, ModelHubClientPool
|
from modelhub_client import ModelHubAPIError, ModelHubClient, ModelHubClientPool
|
||||||
from models import CandidateModel, HFModelSummary, ModelInspection
|
from models import CandidateModel, HFModelSummary, ModelInspection
|
||||||
from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker
|
from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker
|
||||||
|
from submission_claims import DEFAULT_CLAIMS_PATH, SubmissionClaimStore, diversify_candidates
|
||||||
from task_registry import TASK_SPEC_BY_TYPE, all_task_types, choose_framework_for_task, choose_text_generation_framework, pipeline_tags_for_task_types, task_specs_for_model
|
from task_registry import TASK_SPEC_BY_TYPE, all_task_types, choose_framework_for_task, choose_text_generation_framework, pipeline_tags_for_task_types, task_specs_for_model
|
||||||
from template_selector import TemplateSelector
|
from template_selector import TemplateSelector
|
||||||
|
|
||||||
@@ -71,6 +72,11 @@ def build_parser() -> argparse.ArgumentParser:
|
|||||||
parser.add_argument("--runs-dir", default=str(DEFAULT_RUNS_DIR), help=argparse.SUPPRESS)
|
parser.add_argument("--runs-dir", default=str(DEFAULT_RUNS_DIR), help=argparse.SUPPRESS)
|
||||||
parser.add_argument("--ledger-path", default=str(DEFAULT_LEDGER_PATH), help=argparse.SUPPRESS)
|
parser.add_argument("--ledger-path", default=str(DEFAULT_LEDGER_PATH), help=argparse.SUPPRESS)
|
||||||
parser.add_argument("--outcomes-path", default=str(DEFAULT_OUTCOMES_PATH), help=argparse.SUPPRESS)
|
parser.add_argument("--outcomes-path", default=str(DEFAULT_OUTCOMES_PATH), help=argparse.SUPPRESS)
|
||||||
|
parser.add_argument(
|
||||||
|
"--claims-path",
|
||||||
|
default=os.getenv("MODELHUB_AGENT_CLAIMS_PATH", str(DEFAULT_CLAIMS_PATH)),
|
||||||
|
help=argparse.SUPPRESS,
|
||||||
|
)
|
||||||
parser.add_argument("--history-archive-path", default="history/platform_tasks.jsonl", help=argparse.SUPPRESS)
|
parser.add_argument("--history-archive-path", default="history/platform_tasks.jsonl", help=argparse.SUPPRESS)
|
||||||
parser.add_argument("--history-archive-limit", type=int, default=5000, help=argparse.SUPPRESS)
|
parser.add_argument("--history-archive-limit", type=int, default=5000, help=argparse.SUPPRESS)
|
||||||
parser.add_argument("--hf-base-url", default=os.getenv("MODELSCOPE_BASE_URL", "https://modelscope.cn"), help=argparse.SUPPRESS)
|
parser.add_argument("--hf-base-url", default=os.getenv("MODELSCOPE_BASE_URL", "https://modelscope.cn"), help=argparse.SUPPRESS)
|
||||||
@@ -216,9 +222,7 @@ def resolve_max_submit_count(
|
|||||||
remaining_daily_quota: int,
|
remaining_daily_quota: int,
|
||||||
) -> int:
|
) -> int:
|
||||||
explicit_limit = int(getattr(args, "max_submits_per_run", 0) or 0)
|
explicit_limit = int(getattr(args, "max_submits_per_run", 0) or 0)
|
||||||
planned_submit_count = planned_count
|
planned_submit_count = min(planned_count, max(0, remaining_daily_quota))
|
||||||
if remaining_daily_quota > 0:
|
|
||||||
planned_submit_count = min(planned_submit_count, remaining_daily_quota)
|
|
||||||
if explicit_limit > 0:
|
if explicit_limit > 0:
|
||||||
planned_submit_count = min(planned_submit_count, explicit_limit)
|
planned_submit_count = min(planned_submit_count, explicit_limit)
|
||||||
return planned_submit_count
|
return planned_submit_count
|
||||||
@@ -328,16 +332,16 @@ def submit_candidate(
|
|||||||
|
|
||||||
|
|
||||||
def make_run_dir(runs_dir: Path, now) -> Path:
|
def make_run_dir(runs_dir: Path, now) -> Path:
|
||||||
|
runs_dir.mkdir(parents=True, exist_ok=True)
|
||||||
base_name = now.strftime("%Y%m%dT%H%M%SZ")
|
base_name = now.strftime("%Y%m%dT%H%M%SZ")
|
||||||
run_dir = runs_dir / base_name
|
for suffix in range(0, 1000):
|
||||||
if not run_dir.exists():
|
candidate = runs_dir / (base_name if suffix == 0 else f"{base_name}.{suffix:03d}")
|
||||||
return run_dir
|
try:
|
||||||
|
candidate.mkdir(exist_ok=False)
|
||||||
for suffix in range(1, 1000):
|
|
||||||
candidate = runs_dir / f"{base_name}.{suffix:03d}"
|
|
||||||
if not candidate.exists():
|
|
||||||
return candidate
|
return candidate
|
||||||
return candidate
|
except FileExistsError:
|
||||||
|
continue
|
||||||
|
raise RuntimeError(f"Unable to allocate a unique run directory under {runs_dir}")
|
||||||
|
|
||||||
|
|
||||||
def run_submission(
|
def run_submission(
|
||||||
@@ -359,17 +363,17 @@ def run_submission(
|
|||||||
hf_discovery = hf_discovery or HuggingFaceDiscovery(base_url=args.hf_base_url)
|
hf_discovery = hf_discovery or HuggingFaceDiscovery(base_url=args.hf_base_url)
|
||||||
if modelhub_client is None:
|
if modelhub_client is None:
|
||||||
modelhub_tokens = list(getattr(args, "modelhub_tokens", None) or ([] if not args.modelhub_token else [args.modelhub_token]))
|
modelhub_tokens = list(getattr(args, "modelhub_tokens", None) or ([] if not args.modelhub_token else [args.modelhub_token]))
|
||||||
if len(modelhub_tokens) > 1:
|
token_values: list[str | None] = modelhub_tokens or [args.modelhub_token]
|
||||||
clients = [ModelHubClient(token=token, base_url=args.modelhub_base_url) for token in modelhub_tokens]
|
clients = [ModelHubClient(token=token, base_url=args.modelhub_base_url) for token in token_values]
|
||||||
modelhub_client = ModelHubClientPool(clients)
|
modelhub_client = ModelHubClientPool(clients)
|
||||||
else:
|
|
||||||
modelhub_client = ModelHubClient(token=args.modelhub_token, base_url=args.modelhub_base_url)
|
if hasattr(modelhub_client, "begin_cycle"):
|
||||||
|
modelhub_client.begin_cycle()
|
||||||
|
|
||||||
runs_dir = Path(args.runs_dir)
|
runs_dir = Path(args.runs_dir)
|
||||||
ledger_path = Path(args.ledger_path)
|
ledger_path = Path(args.ledger_path)
|
||||||
history_archive_path = Path(args.history_archive_path)
|
history_archive_path = Path(args.history_archive_path)
|
||||||
run_dir = make_run_dir(runs_dir, now)
|
run_dir = make_run_dir(runs_dir, now)
|
||||||
run_dir.mkdir(parents=True, exist_ok=True)
|
|
||||||
ledger_path.parent.mkdir(parents=True, exist_ok=True)
|
ledger_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
history_archive_path.parent.mkdir(parents=True, exist_ok=True)
|
history_archive_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
@@ -387,8 +391,26 @@ def run_submission(
|
|||||||
day_start = now.replace(hour=0, minute=0, second=0, microsecond=0)
|
day_start = now.replace(hour=0, minute=0, second=0, microsecond=0)
|
||||||
# Count today's submissions from the local ledger (avoids expensive paginated API call)
|
# Count today's submissions from the local ledger (avoids expensive paginated API call)
|
||||||
daily_snapshot = count_submissions_for_day(tasks=[], ledger_entries=ledger_entries, day_start=day_start, day_end=now)
|
daily_snapshot = count_submissions_for_day(tasks=[], ledger_entries=ledger_entries, day_start=day_start, day_end=now)
|
||||||
# Fallback: if ledger has no entries yet, do a quick API check (limit to 1 page)
|
# A configured daily target must include every pooled account and every
|
||||||
if daily_snapshot["totalCount"] <= 0:
|
# concurrent worker, not only the first token's first page.
|
||||||
|
if args.daily_target > 0:
|
||||||
|
list_kwargs: dict[str, Any] = {
|
||||||
|
"page_size": 100,
|
||||||
|
"only_mine": True,
|
||||||
|
"begin_time": day_start,
|
||||||
|
"end_time": now,
|
||||||
|
}
|
||||||
|
if isinstance(modelhub_client, ModelHubClientPool):
|
||||||
|
list_kwargs["_fanout_all"] = True
|
||||||
|
today_tasks = modelhub_client.list_tasks(**list_kwargs)
|
||||||
|
daily_snapshot = count_submissions_for_day(
|
||||||
|
tasks=today_tasks,
|
||||||
|
ledger_entries=ledger_entries,
|
||||||
|
day_start=day_start,
|
||||||
|
day_end=now,
|
||||||
|
)
|
||||||
|
# In unlimited mode, keep the inexpensive one-page fallback for an empty ledger.
|
||||||
|
elif daily_snapshot["totalCount"] <= 0:
|
||||||
today_tasks_page = modelhub_client.list_tasks_page(
|
today_tasks_page = modelhub_client.list_tasks_page(
|
||||||
current=1,
|
current=1,
|
||||||
page_size=20,
|
page_size=20,
|
||||||
@@ -533,7 +555,17 @@ def run_submission(
|
|||||||
planned_count=len(candidates),
|
planned_count=len(candidates),
|
||||||
remaining_daily_quota=remaining_daily_quota,
|
remaining_daily_quota=remaining_daily_quota,
|
||||||
)
|
)
|
||||||
planned_candidates = candidates[:planned_submit_count]
|
instance_id = runtime_instance_id()
|
||||||
|
diversified_candidates = diversify_candidates(candidates, instance_id=instance_id)
|
||||||
|
claim_store: SubmissionClaimStore | None = None
|
||||||
|
if args.dry_run:
|
||||||
|
planned_candidates = diversified_candidates[:planned_submit_count]
|
||||||
|
else:
|
||||||
|
claim_store = SubmissionClaimStore(
|
||||||
|
Path(getattr(args, "claims_path", DEFAULT_CLAIMS_PATH)),
|
||||||
|
owner_id=instance_id,
|
||||||
|
)
|
||||||
|
planned_candidates = claim_store.claim(diversified_candidates, limit=planned_submit_count)
|
||||||
submit_workers = 1
|
submit_workers = 1
|
||||||
if not args.dry_run:
|
if not args.dry_run:
|
||||||
submit_workers = resolve_submit_concurrency(
|
submit_workers = resolve_submit_concurrency(
|
||||||
@@ -609,6 +641,12 @@ def run_submission(
|
|||||||
submit_time=result["submitTime"],
|
submit_time=result["submitTime"],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if claim_store is not None:
|
||||||
|
submitted_candidates = [result["candidate"] for result in ordered_results.values() if result.get("outcome") == "submitted"]
|
||||||
|
failed_candidates = [result["candidate"] for result in ordered_results.values() if result.get("outcome") != "submitted"]
|
||||||
|
claim_store.mark_submitted(submitted_candidates)
|
||||||
|
claim_store.release(failed_candidates)
|
||||||
|
|
||||||
write_jsonl(run_dir / "submitted.jsonl", submitted)
|
write_jsonl(run_dir / "submitted.jsonl", submitted)
|
||||||
write_jsonl(run_dir / "skipped.jsonl", skipped)
|
write_jsonl(run_dir / "skipped.jsonl", skipped)
|
||||||
write_jsonl(run_dir / "failed.jsonl", failed)
|
write_jsonl(run_dir / "failed.jsonl", failed)
|
||||||
|
|||||||
@@ -1,13 +1,14 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import hashlib
|
||||||
import os
|
import os
|
||||||
import time
|
|
||||||
import threading
|
import threading
|
||||||
|
import time
|
||||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||||
from datetime import datetime, timedelta
|
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, runtime_instance_id
|
||||||
from defaults import EMBEDDED_MODELHUB_XC_TOKEN
|
from defaults import EMBEDDED_MODELHUB_XC_TOKEN
|
||||||
from http_json import HttpJsonError, JsonHttpClient
|
from http_json import HttpJsonError, JsonHttpClient
|
||||||
|
|
||||||
@@ -233,18 +234,63 @@ def is_active_task(task: dict[str, Any]) -> bool:
|
|||||||
return True
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
CAPACITY_ERROR_MARKERS = (
|
||||||
|
"达到上限",
|
||||||
|
"达上限",
|
||||||
|
"任务数量已达",
|
||||||
|
"队列已满",
|
||||||
|
"queue is full",
|
||||||
|
"queue full",
|
||||||
|
"capacity",
|
||||||
|
"too many active",
|
||||||
|
"active task limit",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def is_capacity_error(error: ModelHubAPIError) -> bool:
|
||||||
|
if error.code in {409, 429}:
|
||||||
|
return True
|
||||||
|
message = str(error).strip().lower()
|
||||||
|
if isinstance(error.payload, dict):
|
||||||
|
message = f"{message} {error.payload.get('message') or ''}".lower()
|
||||||
|
return any(marker in message for marker in CAPACITY_ERROR_MARKERS)
|
||||||
|
|
||||||
|
|
||||||
class ModelHubClientPool:
|
class ModelHubClientPool:
|
||||||
def __init__(self, clients: list[ModelHubClient], *, active_task_cap: int = 100) -> None:
|
def __init__(
|
||||||
|
self,
|
||||||
|
clients: list[ModelHubClient],
|
||||||
|
*,
|
||||||
|
active_task_cap: int | None = None,
|
||||||
|
active_counts_ttl: float | None = None,
|
||||||
|
reservation_ttl: float | None = None,
|
||||||
|
instance_id: str | None = None,
|
||||||
|
) -> None:
|
||||||
if not clients:
|
if not clients:
|
||||||
raise ValueError("At least one ModelHub client is required")
|
raise ValueError("At least one ModelHub client is required")
|
||||||
self.clients = clients
|
self.clients = clients
|
||||||
self.active_task_cap = active_task_cap
|
configured_cap = active_task_cap if active_task_cap is not None else os.getenv("MODELHUB_AGENT_ACTIVE_TASK_CAP", "100")
|
||||||
self._active_counts: list[int] = [0 for _ in clients]
|
self.active_task_cap = max(1, int(configured_cap))
|
||||||
|
configured_ttl = active_counts_ttl if active_counts_ttl is not None else os.getenv("MODELHUB_AGENT_ACTIVE_COUNTS_TTL_SECONDS", "15")
|
||||||
|
configured_reservation_ttl = (
|
||||||
|
reservation_ttl
|
||||||
|
if reservation_ttl is not None
|
||||||
|
else os.getenv("MODELHUB_AGENT_RESERVATION_TTL_SECONDS", "120")
|
||||||
|
)
|
||||||
|
self._active_counts_ttl = max(1.0, float(configured_ttl))
|
||||||
|
self._reservation_ttl = max(self._active_counts_ttl * 2, float(configured_reservation_ttl))
|
||||||
|
self._remote_counts: list[int] = [0 for _ in clients]
|
||||||
self._active_refresh_at: float = 0.0
|
self._active_refresh_at: float = 0.0
|
||||||
self._active_counts_ttl: float = 30.0
|
self._counts_initialized = False
|
||||||
self._state_lock = threading.Lock()
|
self._state_lock = threading.Lock()
|
||||||
# Per-cycle cache for search_by_model_id results (model_id -> merged verify result map)
|
self._refresh_lock = threading.Lock()
|
||||||
self._verify_cache: dict[str, dict[str, Any]] = {}
|
self._reservations: list[dict[int, dict[str, Any]]] = [{} for _ in clients]
|
||||||
|
self._reservation_sequence = 0
|
||||||
|
identity = instance_id or runtime_instance_id()
|
||||||
|
identity_hash = hashlib.sha256(identity.encode("utf-8")).hexdigest()
|
||||||
|
self._selection_cursor = int(identity_hash[:12], 16) % len(clients)
|
||||||
|
self._verify_cache: dict[str, tuple[float, dict[str, Any]]] = {}
|
||||||
|
self._verify_cache_ttl = max(1.0, float(os.getenv("MODELHUB_AGENT_VERIFY_CACHE_TTL_SECONDS", "30")))
|
||||||
# Single reader client to avoid fanout on read operations
|
# Single reader client to avoid fanout on read operations
|
||||||
self._reader = clients[0]
|
self._reader = clients[0]
|
||||||
|
|
||||||
@@ -267,48 +313,86 @@ class ModelHubClientPool:
|
|||||||
except Exception:
|
except Exception:
|
||||||
return []
|
return []
|
||||||
|
|
||||||
|
def _counts_are_fresh_locked(self, now: float) -> bool:
|
||||||
|
return self._counts_initialized and (now - self._active_refresh_at) < self._active_counts_ttl
|
||||||
|
|
||||||
def _refresh_active_counts(self, *, force: bool = False) -> None:
|
def _refresh_active_counts(self, *, force: bool = False) -> None:
|
||||||
now = time.time()
|
now = time.monotonic()
|
||||||
if (not force) and self._active_counts and (now - self._active_refresh_at) < self._active_counts_ttl:
|
with self._state_lock:
|
||||||
return
|
if not force and self._counts_are_fresh_locked(now):
|
||||||
|
return
|
||||||
|
|
||||||
def _to_indexed_result(index: int, client: ModelHubClient) -> tuple[int, int]:
|
# Do not hold the scheduler lock during remote I/O. One refresher is
|
||||||
return index, self._safe_count_active_tasks(client)
|
# enough; all submitting threads can continue using their reservations.
|
||||||
|
with self._refresh_lock:
|
||||||
|
now = time.monotonic()
|
||||||
|
with self._state_lock:
|
||||||
|
if not force and self._counts_are_fresh_locked(now):
|
||||||
|
return
|
||||||
|
|
||||||
results: list[tuple[int, int]] = []
|
def _to_indexed_result(index: int, client: ModelHubClient) -> tuple[int, int]:
|
||||||
with ThreadPoolExecutor(max_workers=min(len(self.clients), 12)) as executor:
|
return index, self._safe_count_active_tasks(client)
|
||||||
futures = {executor.submit(_to_indexed_result, index, client): index for index, client in enumerate(self.clients)}
|
|
||||||
for future in as_completed(futures):
|
|
||||||
index = futures[future]
|
|
||||||
try:
|
|
||||||
results.append(future.result())
|
|
||||||
except Exception:
|
|
||||||
results.append((index, self.active_task_cap))
|
|
||||||
|
|
||||||
self._active_counts = [0 for _ in self.clients]
|
results: list[tuple[int, int]] = []
|
||||||
for index, count in sorted(results, key=lambda item: item[0]):
|
with ThreadPoolExecutor(max_workers=min(len(self.clients), 12)) as executor:
|
||||||
self._active_counts[index] = count
|
futures = {
|
||||||
# Use current time after refresh completes, not the stale 'now' from function start
|
executor.submit(_to_indexed_result, index, client): index
|
||||||
self._active_refresh_at = time.time()
|
for index, client in enumerate(self.clients)
|
||||||
|
}
|
||||||
|
for future in as_completed(futures):
|
||||||
|
index = futures[future]
|
||||||
|
try:
|
||||||
|
results.append(future.result())
|
||||||
|
except Exception:
|
||||||
|
results.append((index, self.active_task_cap))
|
||||||
|
|
||||||
|
refreshed_at = time.monotonic()
|
||||||
|
with self._state_lock:
|
||||||
|
for index, count in results:
|
||||||
|
old_remote_count = self._remote_counts[index]
|
||||||
|
new_remote_count = min(self.active_task_cap, max(0, int(count)))
|
||||||
|
acknowledged = max(0, new_remote_count - old_remote_count)
|
||||||
|
completed_reservations = sorted(
|
||||||
|
(
|
||||||
|
(reservation_id, reservation)
|
||||||
|
for reservation_id, reservation in self._reservations[index].items()
|
||||||
|
if not reservation["inflight"]
|
||||||
|
),
|
||||||
|
key=lambda item: item[1]["updated_at"],
|
||||||
|
)
|
||||||
|
for reservation_id, _reservation in completed_reservations[:acknowledged]:
|
||||||
|
self._reservations[index].pop(reservation_id, None)
|
||||||
|
|
||||||
|
# A remote count can stay flat when one old task finishes as
|
||||||
|
# one new task appears. Expiry prevents that net-zero update
|
||||||
|
# from reserving a slot forever.
|
||||||
|
for reservation_id, reservation in list(self._reservations[index].items()):
|
||||||
|
if reservation["inflight"]:
|
||||||
|
continue
|
||||||
|
if refreshed_at - reservation["updated_at"] >= self._reservation_ttl:
|
||||||
|
self._reservations[index].pop(reservation_id, None)
|
||||||
|
self._remote_counts[index] = new_remote_count
|
||||||
|
self._counts_initialized = True
|
||||||
|
self._active_refresh_at = refreshed_at
|
||||||
|
|
||||||
|
def _effective_count_locked(self, index: int) -> int:
|
||||||
|
return self._remote_counts[index] + len(self._reservations[index])
|
||||||
|
|
||||||
|
def _counts_snapshot_locked(self) -> list[int]:
|
||||||
|
return [min(self.active_task_cap, self._effective_count_locked(index)) for index in range(len(self.clients))]
|
||||||
|
|
||||||
def active_task_counts(self) -> list[int]:
|
def active_task_counts(self) -> list[int]:
|
||||||
|
self._refresh_active_counts()
|
||||||
with self._state_lock:
|
with self._state_lock:
|
||||||
now = time.time()
|
return self._counts_snapshot_locked()
|
||||||
if self._active_counts and (now - self._active_refresh_at) < self._active_counts_ttl:
|
|
||||||
pass # use cached counts
|
|
||||||
else:
|
|
||||||
self._refresh_active_counts()
|
|
||||||
return list(self._active_counts)
|
|
||||||
|
|
||||||
def available_submit_slots(self) -> int:
|
def available_submit_slots(self) -> int:
|
||||||
|
self._refresh_active_counts()
|
||||||
with self._state_lock:
|
with self._state_lock:
|
||||||
# Only refresh if cache is stale (respects TTL) or counts are empty
|
return sum(
|
||||||
now = time.time()
|
max(0, self.active_task_cap - self._effective_count_locked(index))
|
||||||
if self._active_counts and (now - self._active_refresh_at) < self._active_counts_ttl:
|
for index in range(len(self.clients))
|
||||||
pass # use cached counts
|
)
|
||||||
else:
|
|
||||||
self._refresh_active_counts()
|
|
||||||
return sum(max(0, self.active_task_cap - count) for count in self._active_counts)
|
|
||||||
|
|
||||||
def list_tasks(self, **kwargs): # noqa: ANN003, ANN001
|
def list_tasks(self, **kwargs): # noqa: ANN003, ANN001
|
||||||
# Default: use only the reader client to avoid fanout amplification.
|
# Default: use only the reader client to avoid fanout amplification.
|
||||||
@@ -339,10 +423,19 @@ class ModelHubClientPool:
|
|||||||
cached = self._verify_cache.get(model_id)
|
cached = self._verify_cache.get(model_id)
|
||||||
if cached is None:
|
if cached is None:
|
||||||
return None
|
return None
|
||||||
return dict(cached)
|
cached_at, payload = cached
|
||||||
|
if time.monotonic() - cached_at >= self._verify_cache_ttl:
|
||||||
|
self._verify_cache.pop(model_id, None)
|
||||||
|
return None
|
||||||
|
return dict(payload)
|
||||||
|
|
||||||
def _write_to_cache(self, model_id: str, payload: dict[str, Any]) -> None:
|
def _write_to_cache(self, model_id: str, payload: dict[str, Any]) -> None:
|
||||||
self._verify_cache[model_id] = payload
|
self._verify_cache[model_id] = (time.monotonic(), payload)
|
||||||
|
|
||||||
|
def begin_cycle(self) -> None:
|
||||||
|
"""Drop model verification cache entries from the previous scan cycle."""
|
||||||
|
with self._state_lock:
|
||||||
|
self._verify_cache.clear()
|
||||||
|
|
||||||
def search_by_model_id(self, model_id: str) -> dict[str, Any]:
|
def search_by_model_id(self, model_id: str) -> dict[str, Any]:
|
||||||
# Check cache first (per-cycle cache to avoid repeated API calls for the same model)
|
# Check cache first (per-cycle cache to avoid repeated API calls for the same model)
|
||||||
@@ -368,30 +461,87 @@ class ModelHubClientPool:
|
|||||||
def processed_gpus_for_model(self, model_id: str) -> set[str]:
|
def processed_gpus_for_model(self, model_id: str) -> set[str]:
|
||||||
return set(self.get_verify_result_map(model_id).keys())
|
return set(self.get_verify_result_map(model_id).keys())
|
||||||
|
|
||||||
def add_task(self, payload: dict[str, Any]) -> dict[str, Any]:
|
def _reserve_account(self, excluded: set[int]) -> tuple[int, int] | None:
|
||||||
with self._state_lock:
|
with self._state_lock:
|
||||||
self._refresh_active_counts(force=True)
|
remaining_by_index = {
|
||||||
selected_index = None
|
index: self.active_task_cap - self._effective_count_locked(index)
|
||||||
best_remaining = -1
|
for index in range(len(self.clients))
|
||||||
for index, active_count in enumerate(self._active_counts):
|
if index not in excluded
|
||||||
remaining = self.active_task_cap - active_count
|
}
|
||||||
if remaining > best_remaining:
|
usable = {index: remaining for index, remaining in remaining_by_index.items() if remaining > 0}
|
||||||
best_remaining = remaining
|
if not usable:
|
||||||
selected_index = index
|
return None
|
||||||
if selected_index is None or best_remaining <= 0:
|
best_remaining = max(usable.values())
|
||||||
|
tied = [index for index, remaining in usable.items() if remaining == best_remaining]
|
||||||
|
selected_index = min(
|
||||||
|
tied,
|
||||||
|
key=lambda index: (index - self._selection_cursor) % len(self.clients),
|
||||||
|
)
|
||||||
|
self._selection_cursor = (selected_index + 1) % len(self.clients)
|
||||||
|
self._reservation_sequence += 1
|
||||||
|
reservation_id = self._reservation_sequence
|
||||||
|
self._reservations[selected_index][reservation_id] = {
|
||||||
|
"inflight": True,
|
||||||
|
"updated_at": time.monotonic(),
|
||||||
|
}
|
||||||
|
return selected_index, reservation_id
|
||||||
|
|
||||||
|
def _finish_reservation(self, index: int, reservation_id: int, *, succeeded: bool) -> None:
|
||||||
|
with self._state_lock:
|
||||||
|
reservation = self._reservations[index].get(reservation_id)
|
||||||
|
if reservation is None:
|
||||||
|
return
|
||||||
|
if not succeeded:
|
||||||
|
self._reservations[index].pop(reservation_id, None)
|
||||||
|
return
|
||||||
|
reservation["inflight"] = False
|
||||||
|
reservation["updated_at"] = time.monotonic()
|
||||||
|
|
||||||
|
def _mark_account_saturated(self, index: int) -> None:
|
||||||
|
with self._state_lock:
|
||||||
|
self._remote_counts[index] = self.active_task_cap
|
||||||
|
self._counts_initialized = True
|
||||||
|
self._active_refresh_at = time.monotonic()
|
||||||
|
|
||||||
|
def add_task(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||||||
|
self._refresh_active_counts()
|
||||||
|
attempted_accounts: set[int] = set()
|
||||||
|
forced_refresh_done = False
|
||||||
|
last_capacity_error: ModelHubAPIError | None = None
|
||||||
|
|
||||||
|
while True:
|
||||||
|
reservation = self._reserve_account(attempted_accounts)
|
||||||
|
if reservation is None:
|
||||||
|
if not forced_refresh_done:
|
||||||
|
self._refresh_active_counts(force=True)
|
||||||
|
forced_refresh_done = True
|
||||||
|
continue
|
||||||
|
if last_capacity_error is not None:
|
||||||
|
raise last_capacity_error
|
||||||
raise ModelHubAPIError(
|
raise ModelHubAPIError(
|
||||||
f"当前等待中或运行中的异步模型验证任务数量已达上限({self.active_task_cap})"
|
f"当前等待中或运行中的异步模型验证任务数量已达上限({self.active_task_cap})"
|
||||||
)
|
)
|
||||||
self._active_counts[selected_index] += 1
|
|
||||||
|
|
||||||
selected_client = self.clients[selected_index]
|
selected_index, reservation_id = reservation
|
||||||
try:
|
selected_client = self.clients[selected_index]
|
||||||
response = selected_client.add_task(payload)
|
try:
|
||||||
except Exception:
|
response = selected_client.add_task(payload)
|
||||||
with self._state_lock:
|
except ModelHubAPIError as exc:
|
||||||
self._active_counts[selected_index] -= 1
|
self._finish_reservation(selected_index, reservation_id, succeeded=False)
|
||||||
raise
|
if not is_capacity_error(exc):
|
||||||
return response
|
raise
|
||||||
|
# Another process may have filled this account after our count
|
||||||
|
# refresh. Mark it full locally and immediately try another one.
|
||||||
|
self._mark_account_saturated(selected_index)
|
||||||
|
attempted_accounts.add(selected_index)
|
||||||
|
last_capacity_error = exc
|
||||||
|
continue
|
||||||
|
except Exception:
|
||||||
|
self._finish_reservation(selected_index, reservation_id, succeeded=False)
|
||||||
|
raise
|
||||||
|
|
||||||
|
self._finish_reservation(selected_index, reservation_id, succeeded=True)
|
||||||
|
return response
|
||||||
|
|
||||||
def list_tasks_page(self, **kwargs): # noqa: ANN003, ANN001
|
def list_tasks_page(self, **kwargs): # noqa: ANN003, ANN001
|
||||||
"""Single-page task listing via the reader client (no fanout)."""
|
"""Single-page task listing via the reader client (no fanout)."""
|
||||||
|
|||||||
@@ -5,7 +5,7 @@ from datetime import datetime, timedelta
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from common import append_jsonl, parse_datetime, read_jsonl, utc_now, write_jsonl
|
from common import append_jsonl, parse_datetime, read_jsonl, update_jsonl, utc_now
|
||||||
from history_stats import classify_failure, is_failure, is_success
|
from history_stats import classify_failure, is_failure, is_success
|
||||||
from modelhub_client import ModelHubClient, ModelHubClientPool
|
from modelhub_client import ModelHubClient, ModelHubClientPool
|
||||||
|
|
||||||
@@ -25,9 +25,20 @@ class OutcomeTracker:
|
|||||||
self._by_model_gpu: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
|
self._by_model_gpu: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
|
||||||
self._failed_model_gpus: set[tuple[str, str]] = set()
|
self._failed_model_gpus: set[tuple[str, str]] = set()
|
||||||
|
|
||||||
loaded = read_jsonl(self.path)
|
self._records = read_jsonl(self.path)
|
||||||
for record in loaded:
|
self._rebuild_indexes()
|
||||||
self._records.append(record)
|
|
||||||
|
last_sync_times = [
|
||||||
|
parse_datetime(record.get("lastSyncTime"))
|
||||||
|
for record in self._records
|
||||||
|
if record.get("lastSyncTime")
|
||||||
|
]
|
||||||
|
self._last_sync_time: datetime = max(last_sync_times) if last_sync_times else utc_now() - timedelta(days=7)
|
||||||
|
|
||||||
|
def _rebuild_indexes(self) -> None:
|
||||||
|
self._by_task_id.clear()
|
||||||
|
self._by_model_gpu.clear()
|
||||||
|
for record in self._records:
|
||||||
task_id = record.get("taskId")
|
task_id = record.get("taskId")
|
||||||
if task_id:
|
if task_id:
|
||||||
self._by_task_id[str(task_id)] = record
|
self._by_task_id[str(task_id)] = record
|
||||||
@@ -37,13 +48,6 @@ class OutcomeTracker:
|
|||||||
|
|
||||||
self._rebuild_failed_index()
|
self._rebuild_failed_index()
|
||||||
|
|
||||||
last_sync_times = [
|
|
||||||
parse_datetime(record.get("lastSyncTime"))
|
|
||||||
for record in self._records
|
|
||||||
if record.get("lastSyncTime")
|
|
||||||
]
|
|
||||||
self._last_sync_time: datetime = max(last_sync_times) if last_sync_times else utc_now() - timedelta(days=7)
|
|
||||||
|
|
||||||
def record_submission(
|
def record_submission(
|
||||||
self,
|
self,
|
||||||
model_id: str,
|
model_id: str,
|
||||||
@@ -162,8 +166,23 @@ class OutcomeTracker:
|
|||||||
}
|
}
|
||||||
|
|
||||||
def save(self) -> None:
|
def save(self) -> None:
|
||||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
local_records = list(self._records)
|
||||||
write_jsonl(self.path, self._records)
|
|
||||||
|
def merge(existing: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||||
|
merged = list(existing)
|
||||||
|
index_by_key = {_outcome_record_key(record): index for index, record in enumerate(merged)}
|
||||||
|
for record in local_records:
|
||||||
|
key = _outcome_record_key(record)
|
||||||
|
existing_index = index_by_key.get(key)
|
||||||
|
if existing_index is None:
|
||||||
|
index_by_key[key] = len(merged)
|
||||||
|
merged.append(record)
|
||||||
|
continue
|
||||||
|
merged[existing_index] = _prefer_newer_outcome(merged[existing_index], record)
|
||||||
|
return merged
|
||||||
|
|
||||||
|
self._records = update_jsonl(self.path, merge)
|
||||||
|
self._rebuild_indexes()
|
||||||
|
|
||||||
def _rebuild_failed_index(self) -> None:
|
def _rebuild_failed_index(self) -> None:
|
||||||
self._failed_model_gpus.clear()
|
self._failed_model_gpus.clear()
|
||||||
@@ -234,3 +253,31 @@ def _summarize(records: list[dict[str, Any]]) -> dict[str, Any]:
|
|||||||
"pendingRate": round(pending_count / total, 4) if total > 0 else 0.0,
|
"pendingRate": round(pending_count / total, 4) if total > 0 else 0.0,
|
||||||
"failureBreakdown": dict(failure_breakdown),
|
"failureBreakdown": dict(failure_breakdown),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _outcome_record_key(record: dict[str, Any]) -> str:
|
||||||
|
task_id = record.get("taskId")
|
||||||
|
if task_id is not None:
|
||||||
|
return f"task:{task_id}"
|
||||||
|
return "fallback:{model}|{gpu}|{time}".format(
|
||||||
|
model=record.get("modelId") or "",
|
||||||
|
gpu=record.get("targetGpu") or "",
|
||||||
|
time=record.get("submitTime") or "",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _outcome_version(record: dict[str, Any]) -> tuple[int, float, float]:
|
||||||
|
last_sync = parse_datetime(record.get("lastSyncTime"))
|
||||||
|
submit_time = parse_datetime(record.get("submitTime"))
|
||||||
|
outcome_rank = 1 if record.get("outcome") in {"success", "failed"} else 0
|
||||||
|
return (
|
||||||
|
outcome_rank,
|
||||||
|
last_sync.timestamp() if last_sync else 0.0,
|
||||||
|
submit_time.timestamp() if submit_time else 0.0,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _prefer_newer_outcome(existing: dict[str, Any], candidate: dict[str, Any]) -> dict[str, Any]:
|
||||||
|
if _outcome_version(candidate) >= _outcome_version(existing):
|
||||||
|
return candidate
|
||||||
|
return existing
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
import json
|
import json
|
||||||
|
import os
|
||||||
import sys
|
import sys
|
||||||
import time
|
import time
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -10,10 +11,11 @@ from typing import Any, Callable
|
|||||||
from common import utc_now, write_json
|
from common import utc_now, write_json
|
||||||
from daily_runner import DEFAULT_DAILY_RUNS_DIR, log, run_daily_batches
|
from daily_runner import DEFAULT_DAILY_RUNS_DIR, log, run_daily_batches
|
||||||
from hf_discovery import HuggingFaceDiscovery
|
from hf_discovery import HuggingFaceDiscovery
|
||||||
from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR
|
from main import DEFAULT_LEDGER_PATH, DEFAULT_RUNS_DIR, make_run_dir
|
||||||
from modelhub_client import ModelHubClient, ModelHubClientPool
|
from modelhub_client import ModelHubClient, ModelHubClientPool
|
||||||
from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker
|
from outcome_tracker import DEFAULT_OUTCOMES_PATH, OutcomeTracker
|
||||||
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
|
from runner_common import DEFAULT_KEY_PATH, ensure_tokens
|
||||||
|
from submission_claims import DEFAULT_CLAIMS_PATH
|
||||||
from template_selector import TemplateSelector
|
from template_selector import TemplateSelector
|
||||||
|
|
||||||
|
|
||||||
@@ -58,6 +60,11 @@ def build_parser() -> argparse.ArgumentParser:
|
|||||||
parser.add_argument("--key-path", default=str(DEFAULT_KEY_PATH), help="Path to KEY.md containing MODELSCOPE_TOKEN/XC_TOKEN")
|
parser.add_argument("--key-path", default=str(DEFAULT_KEY_PATH), help="Path to KEY.md containing MODELSCOPE_TOKEN/XC_TOKEN")
|
||||||
parser.add_argument("--runs-dir", default=str(DEFAULT_RUNS_DIR), help=argparse.SUPPRESS)
|
parser.add_argument("--runs-dir", default=str(DEFAULT_RUNS_DIR), help=argparse.SUPPRESS)
|
||||||
parser.add_argument("--ledger-path", default=str(DEFAULT_LEDGER_PATH), help=argparse.SUPPRESS)
|
parser.add_argument("--ledger-path", default=str(DEFAULT_LEDGER_PATH), help=argparse.SUPPRESS)
|
||||||
|
parser.add_argument(
|
||||||
|
"--claims-path",
|
||||||
|
default=os.getenv("MODELHUB_AGENT_CLAIMS_PATH", str(DEFAULT_CLAIMS_PATH)),
|
||||||
|
help=argparse.SUPPRESS,
|
||||||
|
)
|
||||||
parser.add_argument("--history-archive-path", default="history/platform_tasks.jsonl", help=argparse.SUPPRESS)
|
parser.add_argument("--history-archive-path", default="history/platform_tasks.jsonl", help=argparse.SUPPRESS)
|
||||||
parser.add_argument("--history-archive-limit", type=int, default=5000, help=argparse.SUPPRESS)
|
parser.add_argument("--history-archive-limit", type=int, default=5000, help=argparse.SUPPRESS)
|
||||||
parser.add_argument("--daily-runs-dir", default=str(DEFAULT_DAILY_RUNS_DIR), help=argparse.SUPPRESS)
|
parser.add_argument("--daily-runs-dir", default=str(DEFAULT_DAILY_RUNS_DIR), help=argparse.SUPPRESS)
|
||||||
@@ -68,9 +75,9 @@ def build_parser() -> argparse.ArgumentParser:
|
|||||||
parser.add_argument("--modelhub-token", default=None, help=argparse.SUPPRESS)
|
parser.add_argument("--modelhub-token", default=None, help=argparse.SUPPRESS)
|
||||||
parser.add_argument("--hf-token", default=None, help=argparse.SUPPRESS)
|
parser.add_argument("--hf-token", default=None, help=argparse.SUPPRESS)
|
||||||
parser.add_argument("--modelscope-token", default=None, help=argparse.SUPPRESS)
|
parser.add_argument("--modelscope-token", default=None, help=argparse.SUPPRESS)
|
||||||
parser.add_argument("--poll-interval-seconds", type=int, default=60, help="Sleep between polling cycles when no slots are available")
|
parser.add_argument("--poll-interval-seconds", type=int, default=15, help="Sleep between polling cycles when no slots are available")
|
||||||
parser.add_argument("--idle-interval-seconds", type=int, default=30, help="Sleep between cycles when a scan submits nothing")
|
parser.add_argument("--idle-interval-seconds", type=int, default=60, help="Sleep between cycles when a scan submits nothing")
|
||||||
parser.add_argument("--post-cycle-cooldown-seconds", type=int, default=0, help="Short sleep after a successful cycle")
|
parser.add_argument("--post-cycle-cooldown-seconds", type=int, default=2, help="Short sleep after a successful cycle")
|
||||||
parser.add_argument("--max-cycles", type=int, default=0, help="Optional hard stop after N cycles; 0 means run until quota is reached")
|
parser.add_argument("--max-cycles", type=int, default=0, help="Optional hard stop after N cycles; 0 means run until quota is reached")
|
||||||
parser.add_argument("--print-stats", action="store_true", help="Load outcomes, sync, print stats report, and exit")
|
parser.add_argument("--print-stats", action="store_true", help="Load outcomes, sync, print stats report, and exit")
|
||||||
return parser
|
return parser
|
||||||
@@ -85,12 +92,11 @@ def _make_cycle_args(base_args: argparse.Namespace) -> argparse.Namespace:
|
|||||||
return cycle_args
|
return cycle_args
|
||||||
|
|
||||||
|
|
||||||
def _build_modelhub_client(base_args: argparse.Namespace) -> ModelHubClient | ModelHubClientPool:
|
def _build_modelhub_client(base_args: argparse.Namespace) -> ModelHubClientPool:
|
||||||
modelhub_tokens = list(getattr(base_args, "modelhub_tokens", None) or ([] if not base_args.modelhub_token else [base_args.modelhub_token]))
|
modelhub_tokens = list(getattr(base_args, "modelhub_tokens", None) or ([] if not base_args.modelhub_token else [base_args.modelhub_token]))
|
||||||
if len(modelhub_tokens) > 1:
|
token_values: list[str | None] = modelhub_tokens or [base_args.modelhub_token]
|
||||||
clients = [ModelHubClient(token=token, base_url=base_args.modelhub_base_url) for token in modelhub_tokens]
|
clients = [ModelHubClient(token=token, base_url=base_args.modelhub_base_url) for token in token_values]
|
||||||
return ModelHubClientPool(clients)
|
return ModelHubClientPool(clients)
|
||||||
return ModelHubClient(token=base_args.modelhub_token, base_url=base_args.modelhub_base_url)
|
|
||||||
|
|
||||||
|
|
||||||
def run_poll_loop(
|
def run_poll_loop(
|
||||||
@@ -110,8 +116,7 @@ def run_poll_loop(
|
|||||||
|
|
||||||
poll_runs_dir = Path(base_args.poll_runs_dir)
|
poll_runs_dir = Path(base_args.poll_runs_dir)
|
||||||
poll_runs_dir.mkdir(parents=True, exist_ok=True)
|
poll_runs_dir.mkdir(parents=True, exist_ok=True)
|
||||||
poll_run_dir = poll_runs_dir / now.strftime("%Y%m%dT%H%M%SZ")
|
poll_run_dir = make_run_dir(poll_runs_dir, now)
|
||||||
poll_run_dir.mkdir(parents=True, exist_ok=True)
|
|
||||||
|
|
||||||
outcome_tracker = outcome_tracker or OutcomeTracker(Path(base_args.outcomes_path))
|
outcome_tracker = outcome_tracker or OutcomeTracker(Path(base_args.outcomes_path))
|
||||||
OUTCOME_SYNC_INTERVAL = 3
|
OUTCOME_SYNC_INTERVAL = 3
|
||||||
|
|||||||
124
modelhub_submmit_api/submission_claims.py
Normal file
124
modelhub_submmit_api/submission_claims.py
Normal file
@@ -0,0 +1,124 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import hashlib
|
||||||
|
from datetime import timedelta
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from common import parse_datetime, runtime_instance_id, update_jsonl, utc_now
|
||||||
|
|
||||||
|
|
||||||
|
DEFAULT_CLAIMS_PATH = Path(".modelhub_state/submission_claims.jsonl")
|
||||||
|
|
||||||
|
|
||||||
|
def candidate_key(candidate: dict[str, Any]) -> str:
|
||||||
|
model_id = candidate.get("repoId") or candidate.get("modelAddress") or "unknown"
|
||||||
|
target_gpu = candidate.get("targetGpu") or "unknown"
|
||||||
|
return f"{model_id}|{target_gpu}"
|
||||||
|
|
||||||
|
|
||||||
|
def diversify_candidates(
|
||||||
|
candidates: list[dict[str, Any]],
|
||||||
|
*,
|
||||||
|
instance_id: str,
|
||||||
|
) -> list[dict[str, Any]]:
|
||||||
|
"""Give concurrent agents different deterministic candidate orders."""
|
||||||
|
|
||||||
|
def sort_key(candidate: dict[str, Any]) -> str:
|
||||||
|
value = f"{instance_id}|{candidate_key(candidate)}"
|
||||||
|
return hashlib.sha256(value.encode("utf-8")).hexdigest()
|
||||||
|
|
||||||
|
return sorted(candidates, key=sort_key)
|
||||||
|
|
||||||
|
|
||||||
|
class SubmissionClaimStore:
|
||||||
|
"""Small filesystem-backed lease store for local multi-process deduplication."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
path: Path | str = DEFAULT_CLAIMS_PATH,
|
||||||
|
*,
|
||||||
|
owner_id: str | None = None,
|
||||||
|
claim_ttl_seconds: int = 600,
|
||||||
|
submitted_ttl_seconds: int = 24 * 60 * 60,
|
||||||
|
) -> None:
|
||||||
|
self.path = Path(path)
|
||||||
|
self.owner_id = owner_id or runtime_instance_id()
|
||||||
|
self.claim_ttl_seconds = max(30, int(claim_ttl_seconds))
|
||||||
|
self.submitted_ttl_seconds = max(self.claim_ttl_seconds, int(submitted_ttl_seconds))
|
||||||
|
|
||||||
|
def claim(
|
||||||
|
self,
|
||||||
|
candidates: list[dict[str, Any]],
|
||||||
|
*,
|
||||||
|
limit: int,
|
||||||
|
) -> list[dict[str, Any]]:
|
||||||
|
if limit <= 0 or not candidates:
|
||||||
|
return []
|
||||||
|
|
||||||
|
selected: list[dict[str, Any]] = []
|
||||||
|
now = utc_now()
|
||||||
|
expires_at = now + timedelta(seconds=self.claim_ttl_seconds)
|
||||||
|
|
||||||
|
def update(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||||
|
active = [record for record in records if not self._is_expired(record, now)]
|
||||||
|
claimed_keys = {str(record.get("key")) for record in active if record.get("key")}
|
||||||
|
for candidate in candidates:
|
||||||
|
key = candidate_key(candidate)
|
||||||
|
if key in claimed_keys:
|
||||||
|
continue
|
||||||
|
active.append(
|
||||||
|
{
|
||||||
|
"key": key,
|
||||||
|
"ownerId": self.owner_id,
|
||||||
|
"state": "claimed",
|
||||||
|
"claimedAt": now.isoformat(),
|
||||||
|
"expiresAt": expires_at.isoformat(),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
claimed_keys.add(key)
|
||||||
|
selected.append(candidate)
|
||||||
|
if len(selected) >= limit:
|
||||||
|
break
|
||||||
|
return active
|
||||||
|
|
||||||
|
update_jsonl(self.path, update)
|
||||||
|
return selected
|
||||||
|
|
||||||
|
def mark_submitted(self, candidates: list[dict[str, Any]]) -> None:
|
||||||
|
keys = {candidate_key(candidate) for candidate in candidates}
|
||||||
|
if not keys:
|
||||||
|
return
|
||||||
|
now = utc_now()
|
||||||
|
expires_at = now + timedelta(seconds=self.submitted_ttl_seconds)
|
||||||
|
|
||||||
|
def update(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||||
|
active = [record for record in records if not self._is_expired(record, now)]
|
||||||
|
for record in active:
|
||||||
|
if record.get("ownerId") == self.owner_id and record.get("key") in keys:
|
||||||
|
record["state"] = "submitted"
|
||||||
|
record["expiresAt"] = expires_at.isoformat()
|
||||||
|
return active
|
||||||
|
|
||||||
|
update_jsonl(self.path, update)
|
||||||
|
|
||||||
|
def release(self, candidates: list[dict[str, Any]]) -> None:
|
||||||
|
keys = {candidate_key(candidate) for candidate in candidates}
|
||||||
|
if not keys:
|
||||||
|
return
|
||||||
|
now = utc_now()
|
||||||
|
|
||||||
|
def update(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||||
|
return [
|
||||||
|
record
|
||||||
|
for record in records
|
||||||
|
if not self._is_expired(record, now)
|
||||||
|
and not (record.get("ownerId") == self.owner_id and record.get("key") in keys)
|
||||||
|
]
|
||||||
|
|
||||||
|
update_jsonl(self.path, update)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _is_expired(record: dict[str, Any], now) -> bool: # noqa: ANN001
|
||||||
|
expires_at = parse_datetime(record.get("expiresAt"))
|
||||||
|
return expires_at is None or expires_at <= now
|
||||||
176
tests/test_concurrency.py
Normal file
176
tests/test_concurrency.py
Normal file
@@ -0,0 +1,176 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
import threading
|
||||||
|
import unittest
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
PACKAGE_DIR = Path(__file__).resolve().parents[1] / "modelhub_submmit_api"
|
||||||
|
if str(PACKAGE_DIR) in sys.path:
|
||||||
|
sys.path.remove(str(PACKAGE_DIR))
|
||||||
|
sys.path.insert(0, str(PACKAGE_DIR))
|
||||||
|
|
||||||
|
from common import read_jsonl # noqa: E402
|
||||||
|
from main import make_run_dir # noqa: E402
|
||||||
|
from modelhub_client import ModelHubAPIError, ModelHubClientPool # noqa: E402
|
||||||
|
from outcome_tracker import OutcomeTracker # noqa: E402
|
||||||
|
from submission_claims import SubmissionClaimStore, candidate_key # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
class FakeClient:
|
||||||
|
def __init__(self, active_count: int = 0, *, reject_capacity: bool = False) -> None:
|
||||||
|
self.active_count = active_count
|
||||||
|
self.reject_capacity = reject_capacity
|
||||||
|
self.submitted: list[dict] = []
|
||||||
|
self.count_calls = 0
|
||||||
|
self.search_calls = 0
|
||||||
|
self._lock = threading.Lock()
|
||||||
|
|
||||||
|
def count_active_tasks(self, **_kwargs) -> int: # noqa: ANN003
|
||||||
|
with self._lock:
|
||||||
|
self.count_calls += 1
|
||||||
|
return self.active_count
|
||||||
|
|
||||||
|
def add_task(self, payload: dict) -> dict:
|
||||||
|
if self.reject_capacity:
|
||||||
|
raise ModelHubAPIError("当前等待中或运行中的异步模型验证任务数量已达上限")
|
||||||
|
with self._lock:
|
||||||
|
self.submitted.append(payload)
|
||||||
|
return {"code": 0, "data": {"id": len(self.submitted)}}
|
||||||
|
|
||||||
|
def search_by_model_id(self, _model_id: str) -> dict:
|
||||||
|
with self._lock:
|
||||||
|
self.search_calls += 1
|
||||||
|
return {"code": 0, "data": {"verifyResult": {}}}
|
||||||
|
|
||||||
|
def list_tasks(self, **_kwargs) -> list[dict]: # noqa: ANN003
|
||||||
|
return []
|
||||||
|
|
||||||
|
def list_tasks_page(self, **_kwargs) -> dict: # noqa: ANN003
|
||||||
|
return {"code": 0, "data": {"records": [], "pages": 0}}
|
||||||
|
|
||||||
|
def find_recent_task_id(self, *_args, **_kwargs): # noqa: ANN002, ANN003
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def make_candidate(index: int) -> dict:
|
||||||
|
return {
|
||||||
|
"repoId": f"owner/model-{index}",
|
||||||
|
"modelAddress": f"https://modelscope.cn/models/owner/model-{index}",
|
||||||
|
"targetGpu": "NVIDIA-A100",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class ClientPoolConcurrencyTests(unittest.TestCase):
|
||||||
|
def test_every_account_reaches_capacity_under_load(self) -> None:
|
||||||
|
clients = [FakeClient() for _ in range(12)]
|
||||||
|
pool = ModelHubClientPool(
|
||||||
|
clients, # type: ignore[arg-type]
|
||||||
|
active_task_cap=10,
|
||||||
|
active_counts_ttl=60,
|
||||||
|
instance_id="full-load-test",
|
||||||
|
)
|
||||||
|
|
||||||
|
with ThreadPoolExecutor(max_workers=24) as executor:
|
||||||
|
list(executor.map(lambda index: pool.add_task({"index": index}), range(120)))
|
||||||
|
|
||||||
|
self.assertEqual([10] * 12, [len(client.submitted) for client in clients])
|
||||||
|
self.assertEqual([10] * 12, pool.active_task_counts())
|
||||||
|
self.assertEqual(0, pool.available_submit_slots())
|
||||||
|
|
||||||
|
def test_concurrent_submissions_reserve_and_balance_accounts(self) -> None:
|
||||||
|
clients = [FakeClient(), FakeClient()]
|
||||||
|
pool = ModelHubClientPool(
|
||||||
|
clients, # type: ignore[arg-type]
|
||||||
|
active_task_cap=4,
|
||||||
|
active_counts_ttl=60,
|
||||||
|
instance_id="balance-test",
|
||||||
|
)
|
||||||
|
|
||||||
|
with ThreadPoolExecutor(max_workers=8) as executor:
|
||||||
|
results = list(executor.map(lambda index: pool.add_task({"index": index}), range(8)))
|
||||||
|
|
||||||
|
self.assertEqual(8, len(results))
|
||||||
|
self.assertEqual([4, 4], [len(client.submitted) for client in clients])
|
||||||
|
self.assertEqual([4, 4], pool.active_task_counts())
|
||||||
|
self.assertEqual(0, pool.available_submit_slots())
|
||||||
|
# A lagging remote count must not erase successful local reservations.
|
||||||
|
pool._refresh_active_counts(force=True)
|
||||||
|
self.assertEqual([4, 4], pool.active_task_counts())
|
||||||
|
|
||||||
|
def test_capacity_rejection_falls_through_to_another_account(self) -> None:
|
||||||
|
full_elsewhere = FakeClient(active_count=0, reject_capacity=True)
|
||||||
|
available = FakeClient(active_count=1)
|
||||||
|
pool = ModelHubClientPool(
|
||||||
|
[full_elsewhere, available], # type: ignore[arg-type]
|
||||||
|
active_task_cap=2,
|
||||||
|
instance_id="fallback-test",
|
||||||
|
)
|
||||||
|
|
||||||
|
response = pool.add_task({"model": "x"})
|
||||||
|
|
||||||
|
self.assertEqual(1, response["data"]["id"])
|
||||||
|
self.assertEqual(0, len(full_elsewhere.submitted))
|
||||||
|
self.assertEqual(1, len(available.submitted))
|
||||||
|
|
||||||
|
def test_verification_cache_is_reset_between_cycles(self) -> None:
|
||||||
|
client = FakeClient()
|
||||||
|
pool = ModelHubClientPool([client], active_task_cap=2) # type: ignore[arg-type]
|
||||||
|
pool.search_by_model_id("owner/model")
|
||||||
|
pool.search_by_model_id("owner/model")
|
||||||
|
self.assertEqual(1, client.search_calls)
|
||||||
|
pool.begin_cycle()
|
||||||
|
pool.search_by_model_id("owner/model")
|
||||||
|
self.assertEqual(2, client.search_calls)
|
||||||
|
|
||||||
|
|
||||||
|
class ProcessCoordinationTests(unittest.TestCase):
|
||||||
|
def test_concurrent_claim_stores_select_disjoint_candidates(self) -> None:
|
||||||
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
||||||
|
path = Path(temporary_dir) / "claims.jsonl"
|
||||||
|
candidates = [make_candidate(index) for index in range(6)]
|
||||||
|
first = SubmissionClaimStore(path, owner_id="worker-a")
|
||||||
|
second = SubmissionClaimStore(path, owner_id="worker-b")
|
||||||
|
|
||||||
|
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||||
|
first_future = executor.submit(first.claim, candidates, limit=3)
|
||||||
|
second_future = executor.submit(second.claim, candidates, limit=3)
|
||||||
|
first_claims = first_future.result()
|
||||||
|
second_claims = second_future.result()
|
||||||
|
|
||||||
|
first_keys = {candidate_key(candidate) for candidate in first_claims}
|
||||||
|
second_keys = {candidate_key(candidate) for candidate in second_claims}
|
||||||
|
self.assertEqual(3, len(first_claims))
|
||||||
|
self.assertEqual(3, len(second_claims))
|
||||||
|
self.assertTrue(first_keys.isdisjoint(second_keys))
|
||||||
|
|
||||||
|
def test_outcome_saves_merge_instead_of_overwriting_other_worker(self) -> None:
|
||||||
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
||||||
|
path = Path(temporary_dir) / "outcomes.jsonl"
|
||||||
|
first = OutcomeTracker(path)
|
||||||
|
second = OutcomeTracker(path)
|
||||||
|
first.record_submission("model-a", "gpu", "vllm", "text-generation", "task-a", "2026-01-01T00:00:00+00:00")
|
||||||
|
second.record_submission("model-b", "gpu", "vllm", "text-generation", "task-b", "2026-01-01T00:00:01+00:00")
|
||||||
|
|
||||||
|
first.save()
|
||||||
|
second.save()
|
||||||
|
|
||||||
|
self.assertEqual({"task-a", "task-b"}, {row["taskId"] for row in read_jsonl(path)})
|
||||||
|
|
||||||
|
def test_run_directories_are_allocated_atomically(self) -> None:
|
||||||
|
with tempfile.TemporaryDirectory() as temporary_dir:
|
||||||
|
base = Path(temporary_dir) / "runs"
|
||||||
|
now = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||||
|
with ThreadPoolExecutor(max_workers=8) as executor:
|
||||||
|
paths = list(executor.map(lambda _index: make_run_dir(base, now), range(8)))
|
||||||
|
|
||||||
|
self.assertEqual(8, len(set(paths)))
|
||||||
|
self.assertTrue(all(path.is_dir() for path in paths))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user