"""リモート実行のサーバ設定(`configs/remote.yaml`)の読み込み・検証と環境変数による上書き."""
from __future__ import annotations
import os
import re
from collections.abc import Mapping
from pathlib import Path
from typing import Any
import yaml
from pydantic import (
BaseModel,
ConfigDict,
Field,
StrictBool,
StrictInt,
field_validator,
model_validator,
)
from tias_judge_model.remote.commands import quote_remote_path, validate_ssh_token
# 手元のプロジェクトルートからの設定ファイルの相対パス
DEFAULT_CONFIG_PATH = Path("configs") / "remote.yaml"
# `jump_host` を上書きする.空文字なら直接接続(研究室 LAN 内にいるとき)
JUMP_ENV_VAR = "TIAS_REMOTE_JUMP"
# 手元の ssh 実行ファイル
SSH_ENV_VAR = "TIAS_REMOTE_SSH"
# `1` でサーバ上実行の検出を無効化する(誤検出時の逃げ道)
ALLOW_SELF_ENV_VAR = "TIAS_REMOTE_ALLOW_SELF"
_SERVER_NAME_RE = re.compile(r"^[A-Za-z0-9_-]+$")
_TRUE_VALUES = frozenset({"1", "true", "yes"})
class _StrictModel(BaseModel):
"""未知のキー(綴り間違い)を拒否し,生成後の変更を禁止する基底モデル."""
model_config = ConfigDict(extra="forbid", frozen=True)
class GpuFreeThreshold(_StrictModel):
"""GPU を「空き」とみなす閾値と,判定前の観測の設定.
Attributes:
max_memory_used_mib: 観測中の使用中メモリ(MiB)の最大がこれ未満.
max_utilization_percent: かつ観測中の使用率(%)の最大がこれ未満なら空き.
observe_sec: この秒数だけ観測してから判定する(0 なら瞬間値のみ).瞬間値だけでは
起動と停止を繰り返す GPU を空きと誤認するため(SPEC_03「判定」).
sample_interval_ms: 観測の間隔(ミリ秒).
"""
max_memory_used_mib: StrictInt = Field(default=1024, ge=0)
max_utilization_percent: StrictInt = Field(default=10, ge=0, le=100)
observe_sec: StrictInt = Field(default=10, ge=0, le=120)
sample_interval_ms: StrictInt = Field(default=500, ge=100, le=5000)
def with_observe_sec(self, observe_sec: int) -> GpuFreeThreshold:
"""観測秒数だけを差し替えた設定を返す(`--observe` 用.範囲の検証も行う).
Args:
observe_sec: 観測秒数.
Returns:
差し替えた設定.
Raises:
pydantic.ValidationError: 範囲外(0〜120 以外)の場合.
"""
return GpuFreeThreshold.model_validate({**self.model_dump(), "observe_sec": observe_sec})
class ServerConfig(_StrictModel):
"""計算サーバ 1 台の設定.
Attributes:
name: サーバ名(英数字・`_`・`-` のみ.`--server` の指定と表示に使う).
address: ssh の接続先アドレス.
user: ssh の接続ユーザー.
priority: GPU 選択の優先度(小さいほど優先.同値は記載順).
remote_dir: サーバ上のプロジェクトルート(`~/...` か絶対パス).
data_dir: サーバ上のデータディレクトリ(ジョブに `TIAS_DATA_ROOT` として渡す).
uv: サーバ上の uv 実行ファイル(非対話 ssh では PATH に無いため絶対パスで呼ぶ).
data_ready: 臨床データが配置済みか(False のサーバは `--no-data` 無しでは選ばない).
"""
name: str
address: str
user: str
priority: StrictInt
remote_dir: str
data_dir: str
uv: str = "~/.local/bin/uv"
data_ready: StrictBool = False
@field_validator("name")
@classmethod
def _check_name(cls, value: str) -> str:
if not _SERVER_NAME_RE.match(value):
raise ValueError(f"サーバ名は英数字・_・- のみ: {value!r}")
return value
@field_validator("address", "user")
@classmethod
def _check_ssh_token(cls, value: str) -> str:
return validate_ssh_token(value, what="接続先")
@field_validator("remote_dir", "data_dir", "uv")
@classmethod
def _check_remote_path(cls, value: str) -> str:
quote_remote_path(value)
return value
class RemoteConfig(_StrictModel):
"""リモート実行の設定のルート.
Attributes:
jump_host: 踏み台(`~/.ssh/config` の Host 名等.`-J` に渡す).None なら直接接続.
connect_timeout_sec: ssh の接続タイムアウト(秒).
gpu_free: 空き判定の閾値.
servers: 計算サーバの一覧(1 件以上.名前は重複不可).
"""
jump_host: str | None = "labpc"
connect_timeout_sec: StrictInt = Field(default=15, ge=1)
gpu_free: GpuFreeThreshold = Field(default_factory=GpuFreeThreshold)
servers: tuple[ServerConfig, ...] = Field(min_length=1)
@field_validator("jump_host")
@classmethod
def _check_jump_host(cls, value: str | None) -> str | None:
return None if value is None else validate_ssh_token(value, what="踏み台")
@model_validator(mode="after")
def _check_unique_names(self) -> RemoteConfig:
# --server の指定は大文字小文字を区別せず照合するため,重複も同じ基準で判定する
seen: set[str] = set()
for server in self.servers:
key = server.name.lower()
if key in seen:
raise ValueError(f"サーバ名が重複している: {server.name!r}")
seen.add(key)
return self
def ordered_servers(self) -> list[ServerConfig]:
"""サーバを優先順(priority 昇順.同値は記載順)に並べて返す.
Returns:
並べ替えたサーバの一覧.
"""
return sorted(self.servers, key=lambda server: server.priority)
def find_server(self, name: str) -> ServerConfig:
"""名前(大文字小文字を区別しない)でサーバを探す.
Args:
name: サーバ名.
Returns:
該当するサーバの設定.
Raises:
ValueError: 該当するサーバが無い場合.
"""
for server in self.servers:
if server.name.lower() == name.lower():
return server
known = ", ".join(server.name for server in self.servers)
raise ValueError(f"サーバ {name!r} は設定に無い(設定済み: {known})")
def load_remote_config(path: Path) -> RemoteConfig:
"""サーバ設定 YAML を読み込み,検証済みの設定を返す.
Args:
path: 設定 YAML のパス.
Returns:
検証済みの設定.
Raises:
ValueError: ファイルが無い,または YAML のトップレベルがマッピングでない場合.
pydantic.ValidationError: 設定値が不正な場合.
"""
if not path.is_file():
raise ValueError(f"サーバ設定ファイルが無い: {path}")
with path.open(encoding="utf-8") as f:
raw: Any = yaml.safe_load(f)
if not isinstance(raw, dict):
raise ValueError(f"サーバ設定 YAML のトップレベルはマッピングでなければならない: {path}")
return RemoteConfig.model_validate(raw)
def resolve_jump_host(config: RemoteConfig, environ: Mapping[str, str]) -> str | None:
"""環境変数 `TIAS_REMOTE_JUMP` を考慮して踏み台を決める.
Args:
config: サーバ設定.
environ: 参照する環境変数.
Returns:
踏み台.None なら直接接続.環境変数が設定されていれば(空文字を含め)それを優先し,
空文字(空白のみを含む)なら直接接続とする.
Raises:
ValueError: 環境変数の値が踏み台として不正な場合.
"""
if JUMP_ENV_VAR not in environ:
return config.jump_host
value = environ[JUMP_ENV_VAR].strip()
return validate_ssh_token(value, what=JUMP_ENV_VAR) if value else None
def resolve_ssh_executable(environ: Mapping[str, str], os_name: str = os.name) -> str:
"""手元の ssh 実行ファイルを決める.
`TIAS_REMOTE_SSH` があればそれを使う.無ければ Windows では System32 の OpenSSH を優先する
(Git Bash から起動すると PATH 上の Git 同梱 ssh(MSYS 版)が先に見つかり,
PowerShell からの起動と挙動が揃わないため).
Args:
environ: 参照する環境変数.
os_name: `os.name` 相当(テスト用に差し替え可能).
Returns:
ssh 実行ファイルのパスまたはコマンド名.
"""
override = environ.get(SSH_ENV_VAR, "").strip()
if override:
return override
if os_name == "nt":
system_root = environ.get("SystemRoot") or environ.get("SYSTEMROOT") or r"C:\Windows"
candidate = Path(system_root) / "System32" / "OpenSSH" / "ssh.exe"
if candidate.is_file():
return str(candidate)
return "ssh"
def is_self_check_disabled(environ: Mapping[str, str]) -> bool:
"""`TIAS_REMOTE_ALLOW_SELF` でサーバ上実行の検出が無効化されているかを返す.
Args:
environ: 参照する環境変数.
Returns:
無効化されていれば True.
"""
return environ.get(ALLOW_SELF_ENV_VAR, "").strip().lower() in _TRUE_VALUES
def detect_self_server(config: RemoteConfig, local_hostname: str) -> ServerConfig | None:
"""手元のホスト名がいずれかのサーバ名と一致するかを判定する.
サーバ上(VS Code Remote の端末等)で実行すると ssh で自分自身へつなぎに行くため,
ssh を起動する前に止める目的.大文字小文字とドメイン部(`newton.local` の `.local`)は
無視して比較する.
Args:
config: サーバ設定.
local_hostname: 手元のホスト名.
Returns:
一致したサーバ.一致しなければ None.
"""
short = local_hostname.strip().lower().split(".", 1)[0]
if not short:
return None
for server in config.servers:
if server.name.lower() == short:
return server
return None