diff --git a/CLAUDE.md b/CLAUDE.md index f68986e..0ce48ec 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -5,7 +5,7 @@ ## 開発進捗 -最新: PLAN_02 Step 1（基盤: 設定スキーマ・manifest・患者単位分割・合成データ生成）完了．次は Step 2（リモート実行ツール） +最新: PLAN_02 Step 2（リモート実行ツール tias-remote．Newton・Archimedes の空き GPU を自動選択）完了．次は Step 3（データセット集計） ※ 本欄は**最新ステップ 1 行のみ上書き更新**．詳細な進捗履歴（動機・設計判断・失敗パターン）は [docs/PROGRESS.md](docs/PROGRESS.md) に追記する．運用ルールは `.claude/rules/progress-log.md` を参照（該当ファイル編集時に自動ロードされる）． ## 必須ルール（コード実装時） @@ -90,6 +90,7 @@ - アーキテクチャ設計: docs/04_SPEC/SPEC_01_アーキテクチャ設計.md - 学習・評価仕様: docs/04_SPEC/SPEC_02_学習・評価仕様.md +- リモート実行仕様: docs/04_SPEC/SPEC_03_リモート実行仕様.md ### 05_TECH（技術設計） diff --git a/configs/remote.yaml b/configs/remote.yaml new file mode 100644 index 0000000..5d8994f --- /dev/null +++ b/configs/remote.yaml @@ -0,0 +1,30 @@ +# リモート実行ツール tias-remote のサーバ設定（SPEC_03「サーバ設定」） +# 臨床データは各サーバの data_dir に事前配置する（ツールはデータを転送しない．ENV_03）． +# 環境変数での上書き: +# TIAS_REMOTE_JUMP 踏み台（jump_host）を上書き．空文字なら直接接続（研究室 LAN 内） +# TIAS_REMOTE_SSH 手元の ssh 実行ファイル +# TIAS_REMOTE_ALLOW_SELF 1 でサーバ上実行の検出を無効化 +jump_host: labpc # 踏み台（~/.ssh/config の Host 名）．null なら直接接続 +connect_timeout_sec: 15 +gpu_free: + max_memory_used_mib: 1024 # 使用中メモリがこれ未満 + max_utilization_percent: 10 # かつ使用率がこれ未満なら「空き」（いずれも観測中の最大値で判定） + observe_sec: 10 # この秒数だけ観測してから判定する（0 なら瞬間値のみ．0〜120） + sample_interval_ms: 500 # 観測の間隔（100〜5000） +servers: + - name: Newton + address: 192.168.23.14 + user: hasegawa + priority: 1 # 小さいほど優先 + remote_dir: ~/work/TIASJudgeModel + data_dir: ~/data/TIASJudgeModel + uv: ~/.local/bin/uv + data_ready: true # 臨床データが配置済みか + - name: Archimedes + address: 192.168.23.15 + user: hasegawa + priority: 2 + remote_dir: ~/work/TIASJudgeModel + data_dir: ~/data/TIASJudgeModel + uv: ~/.local/bin/uv + data_ready: true # 2026-09-27 にデータ配置済み（ENV_03） diff --git "a/docs/02_ENV/ENV_01_\346\212\200\350\241\223\343\202\271\343\202\277\343\203\203\343\202\257.md" "b/docs/02_ENV/ENV_01_\346\212\200\350\241\223\343\202\271\343\202\277\343\203\203\343\202\257.md" index 5310992..0225d88 100644 --- "a/docs/02_ENV/ENV_01_\346\212\200\350\241\223\343\202\271\343\202\277\343\203\203\343\202\257.md" +++ "b/docs/02_ENV/ENV_01_\346\212\200\350\241\223\343\202\271\343\202\277\343\203\203\343\202\257.md" @@ -67,10 +67,12 @@ ## CPU / GPU の依存差 (CPU / GPU Dependencies) -`torch` を PyPI からそのまま依存に入れると，**Linux では CUDA 付きの wheel，Windows では CPU 版の wheel** が選ばれるため，同一の `pyproject.toml` で「手元 = CPU，サーバ = GPU」が成立する見込みである． +**Linux（両サーバ）は PyTorch 公式の CUDA 12.6 版（cu126）**，**Windows（手元）は PyPI の CPU 版**を使う．`pyproject.toml` の `[tool.uv.sources]` で，`torch`・`torchvision` の取得元を `sys_platform == 'linux'` のときだけ cu126 インデックスに向けている． -- サーバでの初回 `uv sync` 後に `torch.cuda.is_available()` が `True` になることを確認する（環境構築手順に含める） -- 想定どおりにならない場合は PyTorch 公式インデックスを `[tool.uv.sources]` で指定する方式に切り替える +- 経緯: 当初は PyPI の wheel をそのまま使い，Linux では CUDA 13 版（cu130）が入っていた．Newton（ドライバ 595・CUDA 13.2 対応）では動いたが，**Archimedes はドライバ 535（CUDA 12.2 まで）のため cu130 版では GPU が使えなかった**（2026-09-27 に判明） +- cu126 版は CUDA のマイナーバージョン互換により 12.2 のドライバでも動く．Archimedes で行列積・畳み込みの逆伝播まで，Newton でも CUDA の動作を確認した（torch は 2.14.0 のまま） +- **両サーバで同じ torch ビルドになる**ため，どちらのサーバで学習しても条件がそろう +- 将来 torch を上げるときは，cu126 版が提供され続けているか（`https://download.pytorch.org/whl/cu126/torch/`）と，Archimedes のドライバ（535）で動くかを確認する ## 採用しなかった選択肢 (Rejected Alternatives) diff --git "a/docs/02_ENV/ENV_02_\347\222\260\345\242\203\346\247\213\347\257\211\346\211\213\351\240\206.md" "b/docs/02_ENV/ENV_02_\347\222\260\345\242\203\346\247\213\347\257\211\346\211\213\351\240\206.md" index 23bb7ce..0b3a296 100644 --- "a/docs/02_ENV/ENV_02_\347\222\260\345\242\203\346\247\213\347\257\211\346\211\213\351\240\206.md" +++ "b/docs/02_ENV/ENV_02_\347\222\260\345\242\203\346\247\213\347\257\211\346\211\213\351\240\206.md" @@ -40,8 +40,12 @@ | Host 名 | アドレス | 用途 | | --- | --- | --- | -| `Newton` | `Newton.local`（研究室 LAN 内） | 学習の実行．RTX A6000×3・Ubuntu 22.04・uv・Python 3.12 | -| `Newton-via-lab` | 研究室 PC を ProxyJump した経路 | 学外から接続する場合 | +| `Newton` | `192.168.23.14`（研究室 LAN 内．`Newton.local`） | 学習の実行．RTX A6000×3・Ubuntu 22.04・uv | +| `Archimedes` | `192.168.23.15`（研究室 LAN 内） | 学習の実行．RTX A6000×2・Ubuntu 22.04・uv（2026-09-27 導入）．**`/home` の空きが少ない**（2026-09-27 時点で使用率 97%・空き 64 GB） | +| 踏み台 `labpc` | 研究室 PC（Tailscale 経由） | 学外から接続する場合は `-J labpc` を付ける | + +- どちらのサーバを使うかはリモート実行ツールが空き GPU から自動で選ぶ（[SPEC_03](../04_SPEC/SPEC_03_リモート実行仕様.md)．Newton 優先） +- `Newton-via-lab` のような別名は `~/.ssh/config` で定義しているが，Archimedes には別名が無い．リモート実行ツールはアドレスと `-J labpc` で直接指定するため別名に依存しない - SSH の鍵と `~/.ssh/config` は **SmTIAS-LightSim で構築済みのものを流用する**（設定手順は同プロジェクトの `docs/02_ENV/ENV_02_環境構築手順.md`「研究室サーバ編」を参照） - **サーバ上では編集しない**．コードは常に手元から送る（次節） @@ -49,12 +53,13 @@ ## 日常の流れ (Daily Workflow) -リモート実行ツール（ENV_01．`sm-remote` の移植）を使う．コマンド名・オプションは実装時に確定し，`ENV_04_開発コマンド.md` に記載する． +リモート実行ツール `tias-remote` を使う（コマンドは [ENV_04](ENV_04_開発コマンド.md)「リモート実行」，仕様は [SPEC_03](../04_SPEC/SPEC_03_リモート実行仕様.md)）． 1. 手元でコードを編集する -2. コードをサーバへ送る（tar over ssh） -3. GPU の空きを確認し，tmux 内で学習を実行する -4. 結果（指標・図・ログ）を回収する +2. `uv run tias-remote gpus` で空き状況を見る（任意） +3. `uv run tias-remote run --name <ジョブ名> -- <コマンド>` で起動する（空いている GPU を両サーバから自動で選び，コードを送ってから tmux で起動する） +4. `uv run tias-remote logs <ジョブ名>` で様子を見る +5. `uv run tias-remote pull <実験ID>` で結果（指標・図・ログ）を持ち帰る ## 環境再現性のルール (Reproducibility Rules) @@ -71,4 +76,5 @@ | 学外からの転送が遅い | Tailscale が DERP 中継経由になっている | `tailscale ping` で経路を確認し，転送は差分だけにする | | ファイル操作系のツールが全てブロックされる | フック用の Python が見つからない | `python3` / `python` / `py` のいずれかが起動できるようにする | | サーバで `uv: command not found` | 非対話 ssh の PATH に `~/.local/bin` が入らない | コマンドの前に `export PATH="$HOME/.local/bin:$PATH"` を付ける（リモート実行ツールもこれを行う） | +| `tias-remote` で「接続不可（Host key verification failed）」 | 手元の `known_hosts` にそのサーバの **IP アドレスのホスト鍵**が無い（別名 `newton.local` 等でしか接続したことが無い） | 一度 `ssh -J labpc hasegawa@ exit` を実行して登録する．既存の別名の鍵と同じかを `ssh-keygen -lF ` と `ssh-keygen -lF <別名>` の指紋で確かめる | | 踏み台経由で `connect failed`（ホスト名不明） | 踏み台側の mDNS（`Newton.local`）の解決が一時的に失敗する | 時間をおいて再試行する．頻発する場合は踏み台の `~/.ssh/config` で Newton を IP 指定にする | diff --git "a/docs/02_ENV/ENV_03_\347\256\241\347\220\206\350\200\205\347\224\250\347\222\260\345\242\203\346\247\213\347\257\211\346\211\213\351\240\206.md" "b/docs/02_ENV/ENV_03_\347\256\241\347\220\206\350\200\205\347\224\250\347\222\260\345\242\203\346\247\213\347\257\211\346\211\213\351\240\206.md" index e1142dc..1443770 100644 --- "a/docs/02_ENV/ENV_03_\347\256\241\347\220\206\350\200\205\347\224\250\347\222\260\345\242\203\346\247\213\347\257\211\346\211\213\351\240\206.md" +++ "b/docs/02_ENV/ENV_03_\347\256\241\347\220\206\350\200\205\347\224\250\347\222\260\345\242\203\346\247\213\347\257\211\346\211\213\351\240\206.md" @@ -47,6 +47,30 @@ `uv sync` は CUDA 版 PyTorch が入るため時間がかかる．`True` と GPU 台数が出ることを確認する．`False` の場合は [ENV_01「CPU / GPU の依存差」](ENV_01_技術スタック.md)に従う． +## Archimedes 側の初期設定 (Archimedes Setup) + +2026-09-27 に実施．Newton と同じ構成にそろえる． + +1. **手元の公開鍵を登録する**（パスワード認証が必要なため人間が実行する） + + ```powershell + type $env:USERPROFILE\.ssh\id_ed25519.pub | ssh -J labpc hasegawa@192.168.23.15 "mkdir -p ~/.ssh && chmod 700 ~/.ssh && cat >> ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys" + ``` + + 聞かれるのは **Archimedes の `hasegawa` アカウントのパスワード**（研究室 PC には鍵で入れるので聞かれない） + +2. **uv をユーザー領域に入れる**（Archimedes には未導入だった）．シェルの設定ファイルを書き換えないよう `UV_NO_MODIFY_PATH=1` を付ける（リモート実行ツールは `~/.local/bin/uv` を絶対パスで呼ぶため PATH の設定は不要） + + ```bash + curl -LsSf https://astral.sh/uv/install.sh | env UV_NO_MODIFY_PATH=1 sh + ``` + +3. ディレクトリを作る（Newton と同じ．データ側は `chmod 700`） +4. 研究室 PC の専用鍵（下記「研究室 PC から Newton への鍵設定」と同じ鍵）を Archimedes の `~/.ssh/authorized_keys` にも追記する +5. 臨床データを研究室 PC から直送する（下記の手順で `$rhost` を `hasegawa@192.168.23.15` にする） + +> ⚠ **Archimedes の `/home` は共用ディスクがほぼ満杯**（2026-09-27 時点で使用率 97%・空き 64 GB）．このプロジェクトの使用量の目安はデータ 0.4 GB・Python 環境（CUDA 版 PyTorch）6〜8 GB・学習結果（重み 1 本約 0.1 GB × 実験数）．重みを溜め込まないこと． + ## 臨床データの初回コピー (Initial Data Copy) 原本は研究室 PC に置き，**研究室 LAN 内で Newton へ一度だけコピーする**．手元 PC を経由させない． @@ -131,7 +155,7 @@ ### 実施記録 (Completion Record) -2026-09-27 に上記手順で転送完了．Newton 側 `~/data/TIASJudgeModel/images/` に **9 ディレクトリ・1085 ファイル（443 MB）**．研究室 PC 側の件数と完全一致し，非 PNG は 0 件．ディレクトリ権限は `drwx------`．転送は研究室 PC → Newton の直送であり，手元 PC は制御コマンドのみを送っている（画像は手元を通っていない）． +2026-09-27 に上記手順で転送完了．**Newton・Archimedes の両方**の `~/data/TIASJudgeModel/images/` に **9 ディレクトリ・1085 ファイル（443 MB）**．研究室 PC 側の件数と完全一致し，非 PNG は 0 件．ディレクトリ権限は `drwx------`．転送は研究室 PC → Newton の直送であり，手元 PC は制御コマンドのみを送っている（画像は手元を通っていない）． 以降は再転送しない．データセットの内容確認（枚数・クラス分布・破損チェック）はサーバ上のコマンドで行い，**集計結果のみ**を手元で見る． diff --git "a/docs/02_ENV/ENV_04_\351\226\213\347\231\272\343\202\263\343\203\236\343\203\263\343\203\211.md" "b/docs/02_ENV/ENV_04_\351\226\213\347\231\272\343\202\263\343\203\236\343\203\263\343\203\211.md" index e8deebd..f70fdcc 100644 --- "a/docs/02_ENV/ENV_04_\351\226\213\347\231\272\343\202\263\343\203\236\343\203\263\343\203\211.md" +++ "b/docs/02_ENV/ENV_04_\351\226\213\347\231\272\343\202\263\343\203\236\343\203\263\343\203\211.md" @@ -2,7 +2,7 @@ 日常の開発で使うコマンドの一覧．**コマンドを推測せず必ず本ファイルを参照する**（新しく確定・変更された時は本ファイルに反映する）． -※ 学習・評価・リモート実行のコマンドは未実装のため本ファイルには未記載（実装時に追記する）． +※ 学習・評価のコマンドは未実装のため本ファイルには未記載（実装時に追記する）． ※ Git Bash ではログの日本語が文字化けする．先頭に `PYTHONIOENCODING=utf-8` を付けると正しく表示される． ## 環境 (Environment) @@ -24,6 +24,31 @@ 設定ファイルは `configs/synthetic.yaml`（手元の合成データ）と `configs/newton.yaml`（サーバの実データ）．環境変数 `TIAS_DATA_ROOT` を設定すると `data.root` より優先される． +## リモート実行 (Remote Execution) + +研究室の GPU サーバ（Newton・Archimedes）で計算するときは `tias-remote` を使う．コマンドはすべて手元で実行する．空いている GPU を両サーバから自動で選ぶ（**未使用の GPU を最優先**．仕様は [SPEC_03](../04_SPEC/SPEC_03_リモート実行仕様.md)）．サーバ設定は `configs/remote.yaml`． + +| 目的 | コマンド | +| --- | --- | +| GPU の状況と `run` が選ぶ GPU | `uv run tias-remote gpus` | +| 実行中ジョブと GPU の状況 | `uv run tias-remote status` | +| コードを送る（全サーバ） | `uv run tias-remote push` | +| 送る一覧だけ見る | `uv run tias-remote push --dry-run` | +| ジョブを起動（GPU 自動選択） | `uv run tias-remote run --name exp1 -- python scripts/xxx.py` | +| GPU を指定して起動 | `uv run tias-remote run --name exp2 --server Newton --gpu 1 -- python ...` | +| 選ばれる GPU だけ確かめる | `uv run tias-remote run --name exp1 --dry-run -- python ...` | +| ログの末尾・追従 | `uv run tias-remote logs exp1`・`uv run tias-remote logs exp1 -f` | +| ジョブを止める | `uv run tias-remote kill exp1` | +| 実験結果を持ち帰る | `uv run tias-remote pull <実験ID>`（重みも要るときは `--with-checkpoints`） | + +- `gpus`・`run`・`status` は GPU を **10 秒観測してから**判定するため，1 回に 13 秒ほどかかる．急ぐときは `--observe 0`（瞬間値のみ．約 2 秒．`gpus`・`run` のみ） +- `run` は起動前に選んだサーバへ自動でコードを送る（`--no-push` で省略）．ジョブには `CUDA_VISIBLE_DEVICES`・`TIAS_DATA_ROOT`（そのサーバのデータの場所）が渡る +- `pull` は[回収ホワイトリスト](../04_SPEC/SPEC_01_アーキテクチャ設計.md)に名前が一致するファイルだけを `runs/<実験ID>/` に持ち帰る．回収しなかったファイルは件数だけ表示される +- ジョブログはサーバ上の `runs/_remote_logs/<ジョブ名>.log`．末尾に `=== [日時] exit <終了コード>` の行が付く +- **研究室 LAN 内にいるとき**は踏み台が不要なので，環境変数 `TIAS_REMOTE_JUMP` を空文字にすると速い（Git Bash: `TIAS_REMOTE_JUMP= uv run tias-remote gpus`．PowerShell: `$env:TIAS_REMOTE_JUMP = ""`） +- **ジョブの実行中に同じサーバへ push しない**（上書き展開なので，実行中のジョブが新旧のコードを混ぜて読むおそれがある．実行中なら警告が出る） +- 各サブコマンドの詳細は `uv run tias-remote <サブコマンド> --help` + ## テスト (Test) | 目的 | コマンド | @@ -56,4 +81,3 @@ - データセットの集計（サーバ上で実行．枚数・患者数・クラス分布・破損チェック） - 学習・評価 -- リモート実行（コード送信・ジョブ実行・結果回収・GPU 空き確認） diff --git "a/docs/03_PLAN/PLAN_02_\351\226\213\347\231\272\343\202\271\343\203\206\343\203\203\343\203\227.md" "b/docs/03_PLAN/PLAN_02_\351\226\213\347\231\272\343\202\271\343\203\206\343\203\203\343\203\227.md" index f8d61b8..d7338de 100644 --- "a/docs/03_PLAN/PLAN_02_\351\226\213\347\231\272\343\202\271\343\203\206\343\203\203\343\203\227.md" +++ "b/docs/03_PLAN/PLAN_02_\351\226\213\347\231\272\343\202\271\343\203\206\343\203\203\343\203\227.md" @@ -7,7 +7,7 @@ | Step | 内容 | 前提 | | --- | --- | --- | | 1 | 基盤: 設定スキーマ・合成データ生成・manifest 読み込み・患者単位分割 ✅ 完了 | なし（手元で完結） | -| 2 | リモート実行ツール（送信・実行・回収・GPU 確認，回収ホワイトリスト） | Step 1 | +| 2 | リモート実行ツール（送信・実行・回収・GPU 確認，回収ホワイトリスト）✅ 完了（Newton・Archimedes の 2 サーバ対応） | Step 1 | | 3 | データセット集計（実データの素性把握） | **臨床データのコピー完了**・倫理確認 | | 4 | 学習パイプライン最小版（ConvNeXt-Tiny を 1 本通す） | Step 2・3 | | 5 | 評価・レポート生成（指標・混同行列・`runs/` 出力） | Step 4 | diff --git "a/docs/04_SPEC/SPEC_01_\343\202\242\343\203\274\343\202\255\343\203\206\343\202\257\343\203\201\343\203\243\350\250\255\350\250\210.md" "b/docs/04_SPEC/SPEC_01_\343\202\242\343\203\274\343\202\255\343\203\206\343\202\257\343\203\201\343\203\243\350\250\255\350\250\210.md" index d7986a6..f2d9dd8 100644 --- "a/docs/04_SPEC/SPEC_01_\343\202\242\343\203\274\343\202\255\343\203\206\343\202\257\343\203\201\343\203\243\350\250\255\350\250\210.md" +++ "b/docs/04_SPEC/SPEC_01_\343\202\242\343\203\274\343\202\255\343\203\206\343\202\257\343\203\201\343\203\243\350\250\255\350\250\210.md" @@ -112,24 +112,32 @@ ├── config.yaml # 解決済み設定 ＋ git コミットハッシュ ＋ seed ├── history.csv # エポックごとの損失・指標 ├── metrics.json # 最終指標（macro F1・accuracy・クラス別） -├── confusion_matrix.png # 図 -├── confusion_matrix.csv # 数値 +├── confusion_matrix.csv # 混同行列の数値 ├── log.txt +├── figures/ # 図専用（学習曲線・混同行列等）．症例画像は書き込まない +│ └── confusion_matrix.png +├── split.json # 分割結果（患者 ID を含む．回収しない） └── checkpoints/best.pt ``` 実験 ID の書式: `YYYYmmdd-HHMMSS_<設定ファイル名>_` +- **`figures/` には図だけを置き，症例画像（舌画像そのもの・切り出し・誤分類例の画像）を書き込まない**．回収ホワイトリストが `figures/*.png` を回収対象にしているため +- 患者 ID・画像ファイル名を含む出力（分割結果・個票の予測）は上記の回収対象外の名前で保存する + ## 回収ホワイトリスト (Pull Whitelist) -手元へ回収できるのは以下のみ．リモート実行ツールにこの制約を実装する． +手元へ回収できるのは以下の**名前が一致するファイルのみ**．列挙外のファイルは回収しない（拡張子や中身での判断はしない）．リモート実行ツール（[SPEC_03](SPEC_03_リモート実行仕様.md)）がこの制約を実装する． -| 回収する | 回収しない | +| 回収する（`runs/<実験ID>/` 配下） | 回収しない | | --- | --- | -| `config.yaml`・`history.csv`・`metrics.json`・`*.png`・集計のみの `*.csv`・`log.txt` | 画像・`manifest.csv`・患者 ID を含むファイル・個票の予測結果 | +| `config.yaml`・`history.csv`・`metrics.json`・`confusion_matrix.csv`・`log.txt` | 上記以外のすべて（`split.json`・個票の予測・画像・manifest 等） | +| `figures/*.png` | `figures/` 以外にある `*.png` | +| `checkpoints/*.pt`（**明示的に指定したときのみ**） | 指定が無いときの `checkpoints/` | -- `checkpoints/` は既定では回収せず，明示的に指定したときのみ回収する -- 個票（画像ごとの予測）が必要な分析はサーバ上で行い，集計結果だけを持ち帰る +- ジョブのログ（`runs/_remote_logs/<ジョブ名>.log`）はファイルとしては回収せず，`tias-remote logs` で表示して確認する（ログに個人情報を出さない規則が前提．`.claude/rules/coding-style.md`） +- データディレクトリ（`~/data/TIASJudgeModel/`）からは何も回収しない +- 個票（画像ごとの予測）が必要な分析はサーバ上で行い，集計結果だけを上記の名前で書き出して持ち帰る ## CLI (Entry Points) @@ -138,6 +146,6 @@ | データセット集計 | 枚数・患者数・クラス分布・画像サイズ・破損チェック | サーバ（結果の要約のみ手元へ） | | 学習 | 設定 YAML を受け取って学習する | サーバ | | 評価 | 学習済み run の評価・レポート生成 | サーバ | -| リモート実行 | コード送信・ジョブ実行・結果回収・GPU 空き確認 | 手元 | +| リモート実行 | コード送信・ジョブ実行・結果回収・GPU 空き確認（複数サーバ．[SPEC_03](SPEC_03_リモート実行仕様.md)） | 手元 | コマンド名は実装時に確定し，[ENV_04 開発コマンド](../02_ENV/ENV_04_開発コマンド.md)に記載する． diff --git "a/docs/04_SPEC/SPEC_03_\343\203\252\343\203\242\343\203\274\343\203\210\345\256\237\350\241\214\344\273\225\346\247\230.md" "b/docs/04_SPEC/SPEC_03_\343\203\252\343\203\242\343\203\274\343\203\210\345\256\237\350\241\214\344\273\225\346\247\230.md" new file mode 100644 index 0000000..47ef57d --- /dev/null +++ "b/docs/04_SPEC/SPEC_03_\343\203\252\343\203\242\343\203\274\343\203\210\345\256\237\350\241\214\344\273\225\346\247\230.md" @@ -0,0 +1,211 @@ +# リモート実行仕様 (Remote Execution) + +手元からコードを送り，研究室の複数の GPU サーバ（Newton・Archimedes）のうち**空いている GPU を自動で選んで**ジョブを実行し，結果を[回収ホワイトリスト](SPEC_01_アーキテクチャ設計.md)に従って持ち帰るツール `tias-remote` の仕様．SmTIAS-LightSim の `sm-remote` を移植し，複数サーバ対応と回収制限を加えたもの． + +## 前提と制約 (Constraints) + +- **臨床データを転送しない**．ツールが扱うのはコードの送信と結果の回収のみ．データは各サーバに事前配置する（[ENV_03](../02_ENV/ENV_03_管理者用環境構築手順.md)） +- サーバに GitHub の認証情報を置かない（コードは常に手元から送る） +- **非対話で動く**（ssh は `BatchMode=yes`．パスワード入力を求めない） +- 共用サーバのため，他のユーザーが使用中の GPU を勝手に使わない + +## サーバ設定 (Server Configuration) + +`configs/remote.yaml` に置き，pydantic で検証する（未知のキーは拒否）． + +```yaml +jump_host: labpc # 踏み台．null なら直接接続 +connect_timeout_sec: 15 +gpu_free: + max_memory_used_mib: 1024 # 使用中メモリがこれ未満 + max_utilization_percent: 10 # かつ使用率がこれ未満なら「空き」 + observe_sec: 10 # この秒数だけ観測してから判定する（0 なら瞬間値のみ） + sample_interval_ms: 500 # 観測の間隔 +servers: + - name: Newton + address: 192.168.23.14 + user: hasegawa + priority: 1 # 小さいほど優先 + remote_dir: ~/work/TIASJudgeModel + data_dir: ~/data/TIASJudgeModel + uv: ~/.local/bin/uv + data_ready: true # 臨床データが配置済みか + - name: Archimedes + address: 192.168.23.15 + user: hasegawa + priority: 2 + remote_dir: ~/work/TIASJudgeModel + data_dir: ~/data/TIASJudgeModel + uv: ~/.local/bin/uv + data_ready: false +``` + +| 検証規則 | 内容 | +| --- | --- | +| `servers` | 1 件以上．`name` は重複不可 | +| `priority` | 整数．同じ値のときは `servers` の記載順で優先 | +| `name` | 英数字・`_`・`-` のみ（コマンド引数・表示に使うため） | +| 閾値 | `max_memory_used_mib` は 0 以上．`max_utilization_percent` は 0〜100 | +| 観測 | `observe_sec` は 0〜120（0 は瞬間値のみで判定）．`sample_interval_ms` は 100〜5000 | +| `connect_timeout_sec` | 1 以上 | + +環境変数による上書き: + +| 環境変数 | 効果 | +| --- | --- | +| `TIAS_REMOTE_JUMP` | `jump_host` を上書きする．**空文字なら直接接続**（研究室 LAN 内にいるとき） | +| `TIAS_REMOTE_SSH` | 手元の ssh 実行ファイル（Windows の既定は `C:\Windows\System32\OpenSSH\ssh.exe`） | +| `TIAS_REMOTE_ALLOW_SELF` | `1` でサーバ上実行の検出を無効化する（誤検出時の逃げ道） | + +## 接続 (Connection) + +- ssh の接続先は `@
`．`jump_host` があれば `-J ` を付ける +- 常に `-o BatchMode=yes -o ConnectTimeout=` を付ける +- サーバ上のコマンドは非対話シェルで動くため，`uv` は設定の絶対パス（`~` 展開）で呼び，`PATH` に `~/.local/bin` を加える（非対話 ssh では `~/.local/bin` が PATH に無い．ENV_02） +- **接続できないサーバは「接続不可」として扱い，他のサーバで処理を続ける**（1 台の不調で全体を止めない）．すべてのサーバに接続できない場合はエラー + +## GPU の空き判定と選択 (GPU Selection) + +### 取得 + +各サーバで次を取得する（ssh 1 回で行う）． + +1. GPU の一覧（1 行 1 GPU） + + ```bash + nvidia-smi --query-gpu=index,uuid,name,memory.used,memory.total,utilization.gpu --format=csv,noheader,nounits + ``` + +2. **観測**: `observe_sec` 秒間，`sample_interval_ms` 間隔で GPU ごとの使用率と使用中メモリを繰り返し取得する（`nvidia-smi --query-gpu=index,utilization.gpu,memory.used --format=csv,noheader,nounits -lms <間隔>` を `timeout <秒>` で打ち切る等） +3. **プロセス別の観測**: 同じ時間，プロセスごとの SM 使用率を取得する（`nvidia-smi pmon -s u -d 1 -c <秒>`．上の観測と並行して実行する） +4. GPU 上の計算プロセス（`gpu_uuid` で GPU と対応付ける） + + ```bash + nvidia-smi --query-compute-apps=gpu_uuid,pid,used_memory --format=csv,noheader,nounits + ``` + +5. 各プロセスの所有ユーザーと**起動からの経過秒数**（`ps -o user=,etimes= -p ` 等） + +※ `--query-compute-apps` と `pmon` の `C`（計算）種別のみを数え，画面描画のプロセス（`G` 種別．Xorg 等）は含めない．`C+G` のように両方を含む種別は計算として数える． +※ nvidia-smi は「最後に使った時刻」の履歴を持たないため取得しない．代わりに観測中の使用率とプロセスの経過時間で判断する． + +### 判定 + +**観測中の最大値**で閾値を判定する（瞬間値だけでは，起動と停止を繰り返す GPU を「空き」と誤認するため．2026-09-27 に Newton で 10 秒間に 0〜74% を往復する GPU を実測）． + +| 状態 | 条件 | 選択 | +| --- | --- | --- | +| **未使用** | 観測中ずっと閾値内（使用率の最大 < `max_utilization_percent` かつ使用中メモリの最大 < `max_memory_used_mib`）で，**計算プロセスが 1 つも無い** | 最優先 | +| **空き（利用者あり）** | 観測中ずっと閾値内だが計算プロセスがある．または計算プロセスの一覧が取得できなかった（利用者不明） | 未使用が無いときのみ | +| 使用中 | 観測中に一度でも閾値を超えた | 選ばない | +| 状態不明 | GPU の行が解釈できない，または観測のサンプルが 1 つも得られない | 選ばない（取得できない GPU を使わないため） | + +- 計算プロセスは**所有者を問わず数える**（自分の別ジョブが載っている GPU に次のジョブを重ねないため） +- 計算プロセスの一覧の取得・解釈に失敗した場合，そのサーバの閾値内の GPU はすべて「空き（利用者あり）」とする（未使用と誤認しないため） +- サーバ単位で GPU の一覧の取得に失敗した場合は，そのサーバの GPU をすべて候補から外す +- `observe_sec: 0`（または `--observe 0`）のときは従来どおり 1 回の瞬間値で判定する + +### 使い方の型 (Usage Pattern) + +表示のために，観測中の使用率から GPU ごとに型を付ける（判定には使わない．判定は上の表）． + +| 型 | 条件 | +| --- | --- | +| 常時使用 | 使用率が 0% を超えたサンプルが全体の 80% 以上 | +| 断続的に使用 | 使用率が 0% を超えたサンプルが 1 つ以上・80% 未満 | +| メモリ確保のみ | 使用率が 0% を超えたサンプルが無く，計算プロセスがある | +| 未使用 | 使用率が 0% を超えたサンプルが無く，計算プロセスも無い | + +プロセス別の観測（pmon）が得られた場合は，プロセスごとにも観測中の SM 使用率の最大値を表示する．得られなかった場合は GPU 単位の型だけを表示する． + +### 選択（`run` で GPU を指定しないとき） + +1. 候補サーバ: `data_ready: true` のサーバ（`--no-data` を付けたときは全サーバ）．`--server` を指定したときはそのサーバのみ +2. **全候補サーバを通して「未使用」の GPU を探し**，`priority` 昇順（同値は記載順）→ GPU の index 昇順で最初のものを選ぶ +3. 未使用が無ければ，同じ順で「空き（利用者あり）」の最初のものを選び，**警告として利用者（ユーザー名）を表示する** +4. どちらも無ければ，各サーバの状況を表示して 0 以外の終了コードで終わる（ジョブは起動しない） + +例: Newton の空き GPU がすべて利用者ありで，Archimedes に未使用の GPU があれば，Newton の方が優先順位が高くても **Archimedes を選ぶ**． + +明示指定: + +- `--server NAME --gpu I` で GPU を直接指定できる．指定した GPU が「未使用」「空き（利用者あり）」なら起動する（後者は利用者を警告表示する）．**使用中・状態不明ならエラー**にし，`--force` を付けたときだけ警告を出して実行する +- `--gpu` だけの指定は不可（どのサーバか曖昧になるため `--server` を必須とする） +- サーバに存在しない GPU 番号を指定した場合はエラー（`--force` でも起動しない） +- `--server` に `data_ready: false` のサーバを指定した場合も，`--no-data` が無ければエラー（候補サーバの規則を明示指定にも適用する） + +### 表示 + +`gpus`・`status` は各 GPU について次を表示する． + +- 状態（未使用・空き（利用者あり）・使用中・状態不明） +- 使い方の型と，観測中の使用率の平均・最大（例: 断続的に使用・平均 21%・最大 74%） +- 計算プロセスごとの所有ユーザー（自分のものは「自分」）・起動からの経過時間（例: 27.7 時間）・観測中の SM 使用率の最大 + +`gpus` と `run` は `--observe SEC` で観測時間を上書きできる（`--observe 0` で瞬間値のみ．急ぐとき用）．観測は全サーバで並行して行う． + +※ 判定から起動までの間に他者が使い始める競合は防げない（既知の制約）． + +## コマンド (Commands) + +| サブコマンド | 動作 | +| --- | --- | +| `gpus [--observe SEC]` | 全サーバの GPU 一覧（状態・使い方の型・利用者）と，`run` が選ぶ GPU を表示する | +| `push [--server NAME]` | コードを送る．省略時は接続できる全サーバ | +| `run --name NAME [--server S] [--gpu I] [--force] [--no-data] [--no-push] [--observe SEC] [--env K=V ...] -- CMD...` | GPU を選び，そのサーバへコードを送り（`--no-push` で省略），tmux でジョブを起動する | +| `status` | 全サーバの実行中ジョブと GPU の状況 | +| `logs NAME [-n N] [-f]` | ジョブのログを表示する．ジョブのあるサーバを自動で探す | +| `kill NAME` | ジョブを止める．ジョブのあるサーバを自動で探す | +| `pull RUN_ID [--server S] [--with-checkpoints]` | 実験結果を回収ホワイトリストに従って持ち帰る | + +### push + +- 送るもの: リポジトリ内のコード一式（`src/`・`scripts/`・`tests/`・`configs/`・`pyproject.toml`・`uv.lock`・`.python-version` 等） +- 送らないもの: `.git/`・`.venv/`・`__pycache__/`・`.pytest_cache/`・`.ruff_cache/`・`data/`・`runs/`・`.claude/` +- 方式: tar over ssh でサーバの `remote_dir` に上書き展開する（Windows に rsync が無いため） +- そのサーバでジョブが実行中なら警告を出す（上書き展開で実行中のジョブが新旧のコードを混ぜて読むおそれがある） + +### run + +- ジョブ名: 英数字で始まり，英数字・`_`・`-` のみ，1〜64 文字 +- **いずれかのサーバで同名のジョブが実行中ならエラー** +- tmux セッション名をジョブ名とし，サーバの `remote_dir` で ` run ` を実行する．接続が切れても継続する +- ジョブに渡す環境変数: `CUDA_VISIBLE_DEVICES=<選んだ GPU の index>`・`TIAS_DATA_ROOT=<そのサーバの data_dir>`・`CUDA_DEVICE_ORDER=PCI_BUS_ID`（CUDA の GPU 番号を nvidia-smi の index にそろえる．そろえないと空きと判定したのと別の GPU を使いうる）・`PYTHONUNBUFFERED=1`・`--env` で指定したもの．**前の 4 つは `--env` で上書きできない**（GPU の割り当てを迂回させないため） +- ログ: `remote_dir/runs/_remote_logs/<ジョブ名>.log` に，開始行（サーバ名・GPU・日時）と終了コード行（`=== [日時] exit <終了コード>`）を付けて追記する +- `--dry-run` を付けると GPU の選択までを行い，送信・起動はしない +- 起動後，選んだサーバと GPU を表示する + +### pull + +- 対象: サーバの `remote_dir/runs//` 配下で，[SPEC_01「回収ホワイトリスト」](SPEC_01_アーキテクチャ設計.md)に**名前が一致するファイルのみ**．`checkpoints/*.pt` は `--with-checkpoints` のときだけ +- `figures/*.png` は `figures/` の**直下のみ**を対象とする（`figures/sub/x.png` は回収しない） +- `RUN_ID` に `/`・`\`・`..` を含むもの，空文字，`.`・`-` で始まるもの，制御文字・前後の空白を含むものは拒否する +- ジョブログ（`runs/_remote_logs/`）は `pull` の対象外．`logs` で表示する +- `--server` を省略したときは全サーバを探し，見つかったのが 1 台ならそこから回収する．複数台にあればエラー（`--server` の指定を求める）．どこにも無ければエラー +- 保存先: 手元の `runs//`（同名ファイルは上書き） +- **二重の防御**: + 1. サーバ側で送るファイルをホワイトリストで絞り込む + 2. 手元で受け取った tar を**展開前に全メンバー検査**し，ホワイトリスト外の名前・パス外に出る名前（`..`・絶対パス・ドライブ文字・バックスラッシュ）・リンク（シンボリック・ハード）が 1 つでもあれば**何も展開せずに中止**する +- 回収しなかったファイルの件数を表示する（名前は表示しない．患者情報を含む名前の可能性があるため） + +### サーバ上での誤実行の防止 + +手元のホスト名がいずれかのサーバの `name` と一致したら，自分自身へ ssh しようとしているとみなして止める（大文字小文字とドメイン部分は無視して比べる．`TIAS_REMOTE_ALLOW_SELF=1` で無効化）． + +### 共通オプションと観測の有無 (Common Options) + +- 全サブコマンドに `--config PATH`（既定はプロジェクトルートの `configs/remote.yaml`） +- 観測（`observe_sec` 秒）を行うのは `gpus`・`run`・`status`．`push`・`kill`・`logs`・`pull` はジョブやファイルの有無だけを見るため観測しない + +## テスト方針 (Testing) + +- ssh・tmux・nvidia-smi はモックする（テストが実サーバに接続しない．`.claude/rules/testing.md`） +- 純粋なロジックとして分離し，単体テストで固定する対象: + - `nvidia-smi` 出力（GPU 状態・観測サンプル・pmon・計算プロセス）と `ps` 出力の解析，`gpu_uuid` による対応付けと 4 段階の判定（**観測中の最大値で判定すること**・状態不明・利用者不明・サンプル 0 件の扱いを含む） + - 使い方の型の判定（常時・断続的・メモリ確保のみ・未使用の境界）と経過時間の表示 + - GPU の選択（**未使用をサーバの優先順位より優先すること**・優先順位・`data_ready`・`--server`・`--gpu`・`--force`・空き無し） + - 接続できないサーバがあっても他のサーバで続行すること + - ssh 引数の組み立て（踏み台の有無・`BatchMode`・タイムアウト） + - 送信対象の選別（除外ディレクトリ） + - **回収ホワイトリストの判定と tar 検査**（最重要．列挙外・パス外・リンクで中止し，何も展開しないこと） + - 設定の検証と環境変数の上書き diff --git a/docs/PROGRESS.md b/docs/PROGRESS.md index 88cc78f..4b62166 100644 --- a/docs/PROGRESS.md +++ b/docs/PROGRESS.md @@ -66,3 +66,15 @@ - 発見 ②: サブエージェントの「ruff 通過」報告が実際と食い違っていた（行長エラー 2 件）．以後，lint・テストの成否は司令塔が自分で実行して確認する - 判断: 画像が 0 枚のクラスディレクトリは当初警告のみで続行していたが，クラスが欠けたまま学習が進むのを防ぐため**エラーで停止する**よう仕様（SPEC_01）と実装を変更した（ユーザー判断）．大文字の拡張子 `.PNG` は受け付けたままとする（実データはすべて小文字で実害なし） +## Step 2: リモート実行ツール + +- [x] **`tias-remote` を実装（Newton・Archimedes の 2 サーバ対応）**: SmTIAS-LightSim の `sm-remote` を移植し，空いている GPU を両サーバから自動で選んでジョブを起動し，結果を回収ホワイトリストに従って持ち帰るツールにした．仕様は [SPEC_03](04_SPEC/SPEC_03_リモート実行仕様.md)． + - 設計判断 ①: GPU の空き判定は **10 秒観測した最大値**で行う．瞬間値だと，起動と停止を繰り返す GPU を空きと誤認する（Newton で 10 秒間に 0〜74% を往復する GPU を実測．瞬間値判定では「空き」と出ていた） + - 設計判断 ②: **計算プロセスが無い GPU（未使用）をサーバの優先順位より優先する**（ユーザー判断）．使用率 0% でも他人のプロセスがメモリを確保している GPU は後回しにする + - 設計判断 ③: 回収ホワイトリストを**名前の完全一致**に厳格化し，図は `figures/` 直下に限定した（`*.png` を広く許すと，誤分類例の舌画像などを持ち帰る経路になる）．pull は受信した tar を全件検査し，1 件でも不正があれば何も展開しない + - 設計判断 ④: ジョブに `CUDA_DEVICE_ORDER=PCI_BUS_ID` を渡す（CUDA と nvidia-smi で GPU 番号がずれ，空きと判定したのとは別の GPU を使う事故を防ぐ） + - 環境: Archimedes をセットアップした（手元の鍵の登録は人間が実施．uv の導入・研究室 PC 専用鍵の登録・臨床データ 1085 枚のコピー・Python 環境）．**Linux の torch を cu126 に統一**した（Archimedes のドライバ 535 は CUDA 12.2 までで，PyPI 既定の cu130 版では GPU が使えなかった．両サーバで同じビルドになる利点もある） + - 失敗パターン: ① IP 指定の ssh が `known_hosts` 未登録で失敗した（これまで別名 `newton.local` でしか接続していなかった）→ IP で登録し，既存の鍵と指紋が一致することを確認 / ② ジョブ終了を待つループが，ログの終了行の書式（実際は `exit N`）を読み違えて空回りした + - 発見: Archimedes の GPU 0 に，42 日間起動したまま観測中ずっと 0% のプロセス（メモリ確保のみ）がある．今回の表示で消し忘れのようなプロセスが一目で分かるようになった + - 未決・既知の制約: 「最後に使った時刻」は nvidia-smi に履歴が無いため取れない（サーバに常駐の記録処理を置く必要があり別作業とした）．計算プロセスの一覧が空で返った場合，本当に 0 件か権限で見えないだけかは区別できない（現状の両サーバでは他人のプロセスが見えている） + diff --git a/pyproject.toml b/pyproject.toml index 4cb1ce0..68e062c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,12 +16,27 @@ "pyyaml>=6.0", ] +[project.scripts] +tias-remote = "tias_judge_model.cli.remote:main" + [dependency-groups] dev = [ "pytest>=8.0", "ruff>=0.8", ] +# Linux（Newton・Archimedes）では torch を PyTorch 公式の CUDA 12.6 版から取る． +# Archimedes のドライバ（535・CUDA 12.2 まで）では PyPI の既定（CUDA 13 版）が動かないため． +# 両サーバで同じビルドにそろう利点もある．Windows（手元）は PyPI の CPU 版のまま（ENV_01）． +[tool.uv.sources] +torch = [{ index = "pytorch-cu126", marker = "sys_platform == 'linux'" }] +torchvision = [{ index = "pytorch-cu126", marker = "sys_platform == 'linux'" }] + +[[tool.uv.index]] +name = "pytorch-cu126" +url = "https://download.pytorch.org/whl/cu126" +explicit = true + [build-system] requires = ["uv_build>=0.12.5,<0.13"] build-backend = "uv_build" diff --git a/src/tias_judge_model/cli/__init__.py b/src/tias_judge_model/cli/__init__.py new file mode 100644 index 0000000..d958497 --- /dev/null +++ b/src/tias_judge_model/cli/__init__.py @@ -0,0 +1 @@ +"""CLI エントリポイント.""" diff --git a/src/tias_judge_model/cli/remote.py b/src/tias_judge_model/cli/remote.py new file mode 100644 index 0000000..4bb4ad7 --- /dev/null +++ b/src/tias_judge_model/cli/remote.py @@ -0,0 +1,1078 @@ +"""tias-remote: 研究室の GPU サーバでの実行を手元から操作する CLI（SPEC_03）. + +手元（Windows の PowerShell / Git Bash）から 1 コマンドで以下を行う．サーバは +`configs/remote.yaml` の一覧（Newton・Archimedes 等）で，空いている GPU を自動で選ぶ． + +- gpus : 全サーバの GPU 一覧と，run が選ぶ GPU を表示する +- status : 全サーバの実行中ジョブと GPU の状況を表示する +- push : コードをサーバへ送る（tar over ssh．rsync 不要） +- run : 空き GPU を選び，そのサーバへコードを送って tmux でジョブを起動する +- logs : ジョブのログを表示する（ジョブのあるサーバを自動で探す） +- kill : ジョブを止める（ジョブのあるサーバを自動で探す） +- pull : 実験結果を回収ホワイトリストに従って持ち帰る + +ssh を呼ぶ副作用はこのモジュールに集め，コマンド文字列の組み立て・出力の解析・GPU の選択・ +回収ホワイトリストの判定は `tias_judge_model.remote` の純関数に任せる． +""" + +from __future__ import annotations + +import argparse +import codecs +import contextlib +import logging +import os +import socket +import subprocess +import sys +import tarfile +import tempfile +import threading +import time +from collections.abc import Callable, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import BinaryIO + +from pydantic import ValidationError + +from tias_judge_model.remote.commands import ( + EXIT_NO_SESSION, + EXIT_NOT_FOUND, + EXIT_REMOTE_DIR_MISSING, + EXIT_SESSION_EXISTS, + EXIT_UNSAFE_PATH, + REMOTE_LOG_DIR, + REMOTE_RUNS_DIR, + build_job_env, + build_job_script, + build_kill_command, + build_log_exists_command, + build_logs_command, + build_probe_command, + build_pull_list_command, + build_pull_tar_command, + build_push_command, + build_run_command, + build_ssh_argv, + job_log_path, + validate_job_name, +) +from tias_judge_model.remote.config import ( + ALLOW_SELF_ENV_VAR, + DEFAULT_CONFIG_PATH, + JUMP_ENV_VAR, + SSH_ENV_VAR, + GpuFreeThreshold, + RemoteConfig, + ServerConfig, + detect_self_server, + is_self_check_disabled, + load_remote_config, + resolve_jump_host, + resolve_ssh_executable, +) +from tias_judge_model.remote.files import PushEntry, collect_push_entries, find_project_root +from tias_judge_model.remote.gpu import ( + ServerProbe, + check_selection_options, + parse_probe_output, + select_gpu, + servers_with_session, +) +from tias_judge_model.remote.report import ( + format_selection, + format_server_report, + format_shared_warning, +) +from tias_judge_model.remote.transfer import copy_stream, extract_pull_archive, write_tar_stream +from tias_judge_model.remote.whitelist import parse_find_listing, plan_pull, validate_run_id +from tias_judge_model.schema import DATA_ROOT_ENV_VAR + +logger = logging.getLogger(__name__) + +# ssh 自身の失敗（接続できない・認証失敗等）の終了コード +_SSH_FAILURE_RC = 255 +# 応答待ちを打ち切ったことを表す内部の終了コード（ssh の終了コードと衝突しない負値） +_TIMEOUT_RC = -1 +# GNU tar の終了コード 1 は「読み取り中にファイルが変化した」等の警告 +_TAR_WARNING_RC = 1 +# 短い問い合わせ（probe・ログの有無・一覧）の応答待ちの上限は接続タイムアウトにこれを足す． +# 踏み台経由では ConnectTimeout が踏み台側の接続に効かないため，全体の上限も設ける +_QUERY_TIMEOUT_MARGIN_SEC = 30 +_MAX_STDERR_CHARS = 160 +_HOST_KEY_FAILURE = b"Host key verification failed" + +_EPILOG = """\ +使用例: + uv run tias-remote gpus # 全サーバの GPU と run が選ぶ GPU + uv run tias-remote status # 実行中ジョブと GPU + uv run tias-remote push # 接続できる全サーバへコードを送る + uv run tias-remote run --name exp1 -- pytest # 空き GPU を自動で選んで起動 + uv run tias-remote run --name exp2 --server Newton --gpu 1 -- python -m ... # GPU を指定 + uv run tias-remote logs exp1 -f # ログを追従 + uv run tias-remote kill exp1 # ジョブを止める + uv run tias-remote pull 20260927-120000_newton_abc1234 # 回収ホワイトリストのみ持ち帰る + +環境変数: + TIAS_REMOTE_JUMP 踏み台を上書き（空文字なら直接接続．研究室 LAN 内にいるとき） + TIAS_REMOTE_SSH 手元の ssh 実行ファイル + TIAS_REMOTE_ALLOW_SELF 1 でサーバ上実行の検出を無効化 +""" + + +@dataclass(frozen=True) +class _Context: + """サブコマンドが使う実行環境. + + Attributes: + config: サーバ設定． + ssh: 手元の ssh 実行ファイル． + jump_host: 踏み台（None なら直接接続）． + project_root: 手元のプロジェクトルート． + """ + + config: RemoteConfig + ssh: str + jump_host: str | None + project_root: Path + + @property + def query_timeout_sec(self) -> int: + """短い問い合わせの応答待ちの上限（秒）.""" + return self.config.connect_timeout_sec + _QUERY_TIMEOUT_MARGIN_SEC + + +@dataclass(frozen=True) +class _SshResult: + """ssh の実行結果. + + Attributes: + rc: 終了コード（応答待ちを打ち切ったら `_TIMEOUT_RC`）． + stdout: 標準出力． + stderr: 標準エラー． + """ + + rc: int + stdout: bytes + stderr: bytes + + +class _CliFormatter(logging.Formatter): + """警告・エラーに日本語の接頭辞を付ける（CLI の利用者向けの表示）.""" + + _PREFIXES = {logging.WARNING: "警告: ", logging.ERROR: "エラー: ", logging.CRITICAL: "エラー: "} + + def format(self, record: logging.LogRecord) -> str: + return self._PREFIXES.get(record.levelno, "") + super().format(record) + + +def _configure_logging() -> None: + """標準エラーへ日本語の接頭辞付きでログを出す（既に設定済みなら何もしない）.""" + handler = logging.StreamHandler(sys.stderr) + handler.setFormatter(_CliFormatter("%(message)s")) + logging.basicConfig(level=logging.INFO, handlers=[handler]) + + +def _write_stdout(text: str) -> None: + """標準出力へ書く（出力先の文字コードで表せない文字は置換する）. + + リモートの出力は UTF-8 で届くが，リダイレクト時の Windows の標準出力は cp932 のことがある． + """ + try: + sys.stdout.write(text) + except UnicodeEncodeError: + encoding = sys.stdout.encoding or "utf-8" + sys.stdout.write(text.encode(encoding, errors="replace").decode(encoding, errors="replace")) + sys.stdout.flush() + + +def _format_size(num_bytes: float) -> str: + """バイト数を読みやすい単位で返す.""" + for unit in ("B", "KB", "MB"): + if num_bytes < 1024: + return f"{num_bytes:.0f} {unit}" if unit == "B" else f"{num_bytes:.1f} {unit}" + num_bytes /= 1024 + return f"{num_bytes:.1f} GB" + + +def _last_stderr_line(stderr: bytes) -> str: + """ssh の標準エラーの最後の空でない行を返す（接続不可の理由の表示用）.""" + lines = [line.strip() for line in stderr.decode("utf-8", errors="replace").splitlines()] + lines = [line for line in lines if line] + return lines[-1][:_MAX_STDERR_CHARS] if lines else "" + + +def _describe_failure(result: _SshResult) -> str: + """ssh が失敗した理由を短く返す.""" + if result.rc == _TIMEOUT_RC: + return "応答がタイムアウトした" + detail = _last_stderr_line(result.stderr) if result.rc == _SSH_FAILURE_RC else "" + return f"終了コード {result.rc}" + (f": {detail}" if detail else "") + + +# --------------------------------------------------------------------------- +# ssh の実行（副作用） +# --------------------------------------------------------------------------- + + +def _ssh_argv( + ctx: _Context, server: ServerConfig, command: str, *, stdin: bool = False +) -> list[str]: + """サーバでリモートコマンドを実行する ssh の引数リストを返す.""" + return build_ssh_argv( + ctx.ssh, + user=server.user, + address=server.address, + jump_host=ctx.jump_host, + connect_timeout_sec=ctx.config.connect_timeout_sec, + remote_command=command, + stdin=stdin, + ) + + +def _popen(argv: list[str], **kwargs: object) -> subprocess.Popen[bytes]: + """ssh を起動する（見つからなければ対処を添えて ValueError にする）.""" + try: + return subprocess.Popen(argv, **kwargs) # noqa: S603 + except FileNotFoundError as exc: + raise ValueError( + f"ssh が見つからない: {argv[0]}（環境変数 {SSH_ENV_VAR} で実行ファイルを指定できる）" + ) from exc + + +def _run_ssh( + argv: list[str], *, input_bytes: bytes | None = None, timeout: float | None = None +) -> _SshResult: + """ssh を実行して終了を待ち，標準出力・標準エラーを受け取る. + + Args: + argv: ssh の引数リスト． + input_bytes: 標準入力へ渡すバイト列（None なら何も渡さない）． + timeout: 応答待ちの上限（秒）．超えたら ssh を止めて `_TIMEOUT_RC` を返す． + + Returns: + 実行結果． + """ + stdin = subprocess.PIPE if input_bytes is not None else subprocess.DEVNULL + proc = _popen(argv, stdin=stdin, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + try: + out, err = proc.communicate(input=input_bytes, timeout=timeout) + except subprocess.TimeoutExpired: + proc.kill() + out, err = proc.communicate() + return _SshResult(_TIMEOUT_RC, out or b"", err or b"") + return _SshResult(proc.returncode, out or b"", err or b"") + + +def _map_servers[T](servers: Sequence[ServerConfig], func: Callable[[ServerConfig], T]) -> list[T]: + """サーバごとの処理を並列に実行し，servers の順で結果を返す. + + 接続できないサーバの接続タイムアウトを待つ間，他のサーバの処理を止めないため並列にする． + """ + if not servers: + return [] + with ThreadPoolExecutor(max_workers=len(servers)) as executor: + return list(executor.map(func, servers)) + + +def _query_servers( + ctx: _Context, servers: Sequence[ServerConfig], build_command: Callable[[ServerConfig], str] +) -> list[_SshResult]: + """短い問い合わせを各サーバで並列に実行し，servers の順で結果を返す. + + コマンドは ssh を起動する前にすべて組み立てる（引数の検証エラーで一部のサーバにだけ + 問い合わせた状態にしないため）． + """ + commands = {server.name: build_command(server) for server in servers} + return _map_servers( + servers, + lambda s: _run_ssh(_ssh_argv(ctx, s, commands[s.name]), timeout=ctx.query_timeout_sec), + ) + + +def _as_exit_code(rc: int) -> int: + """ssh の結果の終了コードを CLI の終了コードにする（応答待ちの打ち切り等の負値は 1）.""" + return rc if rc >= 0 else 1 + + +def _ssh_hint(ctx: _Context, server: ServerConfig, result: _SshResult) -> str: + """よくある ssh の失敗に対処の案内を返す（該当しなければ空文字）.""" + if result.rc != _SSH_FAILURE_RC or _HOST_KEY_FAILURE not in result.stderr: + return "" + # BatchMode ではホスト鍵の確認に答えられないため，初回だけ手で登録してもらう + jump = f"-J {ctx.jump_host} " if ctx.jump_host else "" + return ( + f" → ホスト鍵が未登録: 手元で一度 ssh {jump}{server.user}@{server.address} exit を" + "実行して登録する" + ) + + +def _probe_server(ctx: _Context, server: ServerConfig, threshold: GpuFreeThreshold) -> ServerProbe: + """サーバ 1 台の tmux セッションと GPU の状況を取得する（接続できなければその旨を返す）. + + GPU の観測はサーバ側で `threshold.observe_sec` 秒かかるため，応答待ちの上限をそのぶん延ばす． + """ + command = build_probe_command( + server.remote_dir, + observe_sec=threshold.observe_sec, + sample_interval_ms=threshold.sample_interval_ms, + ) + argv = _ssh_argv(ctx, server, command) + result = _run_ssh(argv, timeout=ctx.query_timeout_sec + threshold.observe_sec) + if result.rc != 0: + return ServerProbe(server, None, _describe_failure(result) + _ssh_hint(ctx, server, result)) + try: + output = parse_probe_output(result.stdout.decode("utf-8", errors="replace"), threshold) + except ValueError as exc: + return ServerProbe(server, None, str(exc)) + return ServerProbe(server, output) + + +def _probe_servers( + ctx: _Context, servers: Sequence[ServerConfig], threshold: GpuFreeThreshold | None = None +) -> list[ServerProbe]: + """複数サーバの状況を並列に取得する（servers の順で返す）. + + Args: + ctx: 実行環境． + servers: 対象のサーバ． + threshold: 判定と観測の設定．None なら観測しない（ジョブの有無だけが要る push・kill + 用．瞬間値で判定する）． + """ + if threshold is None: + threshold = ctx.config.gpu_free.with_observe_sec(0) + elif threshold.observe_sec > 0: + logger.info( + "GPU を %d 秒観測している（%d ミリ秒間隔．全サーバ並行）", + threshold.observe_sec, + threshold.sample_interval_ms, + ) + return _map_servers(servers, lambda server: _probe_server(ctx, server, threshold)) + + +def _threshold_for(ctx: _Context, observe_sec: int | None) -> GpuFreeThreshold: + """`--observe` があれば観測秒数を差し替えた判定の設定を返す（範囲外は ValueError）.""" + threshold = ctx.config.gpu_free + if observe_sec is None: + return threshold + try: + return threshold.with_observe_sec(observe_sec) + except ValidationError as exc: + # 設定ファイルではなくコマンドライン引数の誤りとして伝える + raise ValueError(f"--observe は 0〜120 の整数で指定する（指定 {observe_sec}）") from exc + + +def _target_servers(ctx: _Context, server_name: str | None) -> list[ServerConfig]: + """`--server` の指定があればそのサーバ，無ければ全サーバ（優先順）を返す.""" + if server_name is not None: + return [ctx.config.find_server(server_name)] + return ctx.config.ordered_servers() + + +def _warn_unreachable(probes: Sequence[ServerProbe]) -> None: + """接続できなかったサーバを警告し，すべて接続できなければ ValueError にする.""" + reachable = [probe for probe in probes if probe.reachable] + if not reachable: + details = "; ".join(f"{p.server.name}: {p.error}" for p in probes) + raise ValueError(f"どのサーバにも接続できない（{details}）") + for probe in probes: + if not probe.reachable: + logger.warning( + "%s に接続できない（%s）．他のサーバで続行する", probe.server.name, probe.error + ) + + +def _report_remote_failure(result: _SshResult, server: ServerConfig) -> None: + """ssh・リモートスクリプトに共通の失敗理由をログに出す.""" + if result.rc == EXIT_REMOTE_DIR_MISSING: + logger.error("%s に %s が無い．先に tias-remote push する", server.name, server.remote_dir) + elif result.rc == _SSH_FAILURE_RC: + logger.error( + "%s へ ssh 接続できなかった（%s）", server.name, _last_stderr_line(result.stderr) + ) + else: + logger.error( + "%s でリモートコマンドが失敗した（%s）", server.name, _describe_failure(result) + ) + + +def _pick_single(found: Sequence[ServerConfig], *, what: str, unreachable: int) -> ServerConfig: + """探索で見つかったサーバが 1 台であることを確かめて返す.""" + if len(found) > 1: + names = ", ".join(server.name for server in found) + raise ValueError(f"{what}が複数のサーバ（{names}）にある．--server で指定する") + if not found: + note = f"（接続できないサーバが {unreachable} 台ある）" if unreachable else "" + raise ValueError(f"{what}がどのサーバにも見つからない{note}") + return found[0] + + +# --------------------------------------------------------------------------- +# gpus / status +# --------------------------------------------------------------------------- + + +def _cmd_gpus(args: argparse.Namespace, ctx: _Context) -> int: + """gpus サブコマンド: 全サーバの GPU 一覧と run が選ぶ GPU を表示する.""" + threshold = _threshold_for(ctx, args.observe) + probes = _probe_servers(ctx, ctx.config.ordered_servers(), threshold) + _write_stdout(format_server_report(probes, show_jobs=False)) + if not any(probe.reachable for probe in probes): + raise ValueError("どのサーバにも接続できない") + selection = select_gpu(probes, include_not_ready=args.no_data) + _write_stdout(format_selection(selection) + "\n") + if selection is None and not args.no_data: + alternative = select_gpu(probes, include_not_ready=True) + if alternative is not None: + _write_stdout( + f" （データ未配置のサーバも含めると {alternative.server.name} の GPU " + f"{alternative.gpu_index}．データを使わないジョブなら --no-data で選べる）\n" + ) + return 0 + + +def _cmd_status(args: argparse.Namespace, ctx: _Context) -> int: + """status サブコマンド: 全サーバの実行中ジョブと GPU の状況を表示する.""" + probes = _probe_servers(ctx, ctx.config.ordered_servers(), ctx.config.gpu_free) + _write_stdout(format_server_report(probes, show_jobs=True)) + if not any(probe.reachable for probe in probes): + raise ValueError("どのサーバにも接続できない") + return 0 + + +# --------------------------------------------------------------------------- +# push +# --------------------------------------------------------------------------- + + +def _stream_push(argv: list[str], entries: Sequence[PushEntry]) -> int: + """tar ストリームを ssh の標準入力へ流し，ssh の終了コードを返す.""" + proc = _popen(argv, stdin=subprocess.PIPE) + stdin = proc.stdin + if stdin is None: + raise RuntimeError("ssh の標準入力を開けなかった") + try: + write_tar_stream(entries, stdin) + except OSError as exc: + # ssh が先に終了してパイプが閉じている．後始末の close 失敗は原因ではないので無視する + with contextlib.suppress(OSError): + stdin.close() + rc = proc.wait() + logger.error("送信が中断された（ssh 終了コード %d）: %s", rc, exc) + return rc or 1 + stdin.close() + return proc.wait() + + +def _push_to(ctx: _Context, probe: ServerProbe, entries: Sequence[PushEntry]) -> int: + """サーバ 1 台へコードを送り，終了コードを返す.""" + server = probe.server + if probe.output is not None and probe.output.jobs: + logger.warning( + "%s では実行中のジョブがある（%s）．上書き展開のため，実行中のジョブが" + "新旧のコードを混ぜて読むおそれがある", + server.name, + ", ".join(probe.output.jobs), + ) + total = sum(entry.size for entry in entries) + logger.info( + "送信: %d ファイル / %s → %s:%s", + len(entries), + _format_size(total), + server.name, + server.remote_dir, + ) + start = time.monotonic() + argv = _ssh_argv(ctx, server, build_push_command(server.remote_dir), stdin=True) + rc = _stream_push(argv, entries) + if rc != 0: + logger.error("%s への送信に失敗した（終了コード %d）", server.name, rc) + return rc + logger.info("送信完了: %s（%.1f 秒）", server.name, time.monotonic() - start) + return 0 + + +def _collect_entries(ctx: _Context) -> list[PushEntry]: + """送信対象を集める（空ならエラー）.""" + entries = collect_push_entries(ctx.project_root) + if not entries: + raise ValueError("送信するファイルが無い") + return entries + + +def _cmd_push(args: argparse.Namespace, ctx: _Context) -> int: + """push サブコマンド: コードをサーバへ送る（省略時は接続できる全サーバ）.""" + entries = _collect_entries(ctx) + servers = _target_servers(ctx, args.server) + if args.dry_run: + for entry in entries: + _write_stdout(f" {_format_size(entry.size):>10} {entry.arcname}\n") + names = ", ".join(server.name for server in servers) + _write_stdout( + f"送信対象: {len(entries)} ファイル / " + f"{_format_size(sum(e.size for e in entries))}（送信先の候補: {names}）\n" + "（--dry-run のため送信していない）\n" + ) + return 0 + + probes = _probe_servers(ctx, servers) + _warn_unreachable(probes) + rc_total = 0 + for probe in probes: + if probe.reachable: + rc = _push_to(ctx, probe, entries) + rc_total = rc_total or rc + return rc_total + + +# --------------------------------------------------------------------------- +# run +# --------------------------------------------------------------------------- + + +def _cmd_run(args: argparse.Namespace, ctx: _Context) -> int: + """run サブコマンド: 空き GPU を選び，コードを送って tmux でジョブを起動する.""" + command = list(args.command) + if command and command[0] == "--": + command = command[1:] + if not command: + raise ValueError( + "実行するコマンドを -- の後に指定する（例: tias-remote run --name x -- python ...）" + ) + name = validate_job_name(args.name) + check_selection_options(server_name=args.server, gpu_index=args.gpu, force=args.force) + if args.server is not None: + ctx.config.find_server(args.server) + threshold = _threshold_for(ctx, args.observe) + + probes = _probe_servers(ctx, ctx.config.ordered_servers(), threshold) + _warn_unreachable(probes) + running = servers_with_session(probes, name, jobs_only=False) + if running: + names = ", ".join(probe.server.name for probe in running) + raise ValueError( + f"ジョブ名 {name!r} の tmux セッションが {names} で実行中である" + "（--name で別名を付けるか，終了を待つ）" + ) + if any(not probe.reachable for probe in probes): + logger.warning("接続できないサーバでは同名ジョブの有無を確認していない") + + selection = select_gpu( + probes, + server_name=args.server, + gpu_index=args.gpu, + force=args.force, + include_not_ready=args.no_data, + ) + if selection is None: + _write_stdout(format_server_report(probes, show_jobs=True)) + scope = "全サーバ" if args.no_data else "データ配置済みのサーバ" + logger.error( + "%s（%s）に未使用・空き（利用者あり）の GPU が無いため，ジョブを起動しなかった", + args.server or scope, + "--no-data 指定" if args.no_data else "data_ready: true", + ) + return 1 + if selection.forced: + logger.warning( + "%s の GPU %d は%sだが，--force により使う（他のユーザーの計算と競合しうる）", + selection.server.name, + selection.gpu_index, + selection.state, + ) + shared_warning = format_shared_warning(selection) + if shared_warning is not None: + logger.warning("%s", shared_warning) + + server = selection.server + env = build_job_env(selection.gpu_index, args.env) + script = build_job_script( + server.remote_dir, + name, + command, + uv_path=server.uv, + server_name=server.name, + gpu_index=selection.gpu_index, + env=env, + path_env={DATA_ROOT_ENV_VAR: server.data_dir}, + ) + remote_command = build_run_command(server.remote_dir, name, script) + if args.dry_run: + _write_stdout( + f"選んだ GPU: {server.name} の GPU {selection.gpu_index}（{selection.state.value}）\n" + f"実行するリモートコマンド:\n{remote_command}\n" + "（--dry-run のため送信・起動していない）\n" + ) + return 0 + + if not args.no_push: + probe = next(p for p in probes if p.server.name == server.name) + rc = _push_to(ctx, probe, _collect_entries(ctx)) + if rc != 0: + return rc + + result = _run_ssh(_ssh_argv(ctx, server, remote_command), timeout=ctx.query_timeout_sec) + if result.rc == EXIT_SESSION_EXISTS: + logger.error( + "%s で同名のジョブ（tmux セッション）%r が実行中である（--name で別名を付ける）", + server.name, + name, + ) + return result.rc + if result.rc != 0: + _report_remote_failure(result, server) + return _as_exit_code(result.rc) + + log_path = f"{server.remote_dir}/{job_log_path(name)}" + _write_stdout( + f"ジョブを起動した: {name}（{server.name} の GPU {selection.gpu_index}）\n" + f" 環境変数: {' '.join(f'{k}={v}' for k, v in env.items())} " + f"{DATA_ROOT_ENV_VAR}={server.data_dir}\n" + f" ログ : {server.name}:{log_path}\n" + " 状態確認: uv run tias-remote status\n" + f" ログ追従: uv run tias-remote logs {name} -f\n" + ) + return 0 + + +# --------------------------------------------------------------------------- +# logs / kill +# --------------------------------------------------------------------------- + + +def _locate_log(ctx: _Context, name: str, server_name: str | None) -> ServerConfig: + """ジョブログがあるサーバを探す（複数・0 台ならエラー）.""" + servers = _target_servers(ctx, server_name) + results = _query_servers(ctx, servers, lambda s: build_log_exists_command(s.remote_dir, name)) + found: list[ServerConfig] = [] + unreachable = 0 + for server, result in zip(servers, results, strict=True): + if result.rc == 0: + found.append(server) + elif result.rc != EXIT_NOT_FOUND: + unreachable += 1 + logger.warning("%s に接続できない（%s）", server.name, _describe_failure(result)) + return _pick_single(found, what=f"ジョブ {name!r} のログ", unreachable=unreachable) + + +def _follow(argv: list[str]) -> int: + """ssh の出力を逐次表示し続ける（Ctrl+C で終了）. + + 進捗バーは改行せず \\r で上書きすることがあるため，行単位ではなく受け取った分から表示する． + ssh の標準入力はパイプにして開いたままにし，終了時に閉じてリモートの tail を止める． + """ + proc = _popen(argv, stdin=subprocess.PIPE, stdout=subprocess.PIPE) + if proc.stdin is None or proc.stdout is None: + raise RuntimeError("ssh の標準入出力を開けなかった") + decoder = codecs.getincrementaldecoder("utf-8")(errors="replace") + interrupted = False + try: + while chunk := proc.stdout.read1(65536): + _write_stdout(decoder.decode(chunk)) + except KeyboardInterrupt: + interrupted = True + logger.info("（追従を終了した．ジョブは継続している）") + finally: + # stdin を閉じるとリモートの cat が EOF を受けて tail を止める．既に ssh が + # 終了していれば close は失敗しうるが，後始末なので無視する + with contextlib.suppress(OSError): + proc.stdin.close() + rc = proc.wait() + # Ctrl+C は追従の正常な終わらせ方（Windows では ssh も同時に止まり非 0 を返す） + return 0 if interrupted else rc + + +def _cmd_logs(args: argparse.Namespace, ctx: _Context) -> int: + """logs サブコマンド: ジョブのログを表示する（ジョブのあるサーバを自動で探す）.""" + name = validate_job_name(args.name) + if args.lines < 0: + raise ValueError("--lines は 0 以上を指定する") + server = _locate_log(ctx, name, args.server) + command = build_logs_command(server.remote_dir, name, lines=args.lines, follow=args.follow) + if args.follow: + logger.info("%s のジョブ %s のログを追従する（Ctrl+C で終了）", server.name, name) + return _follow(_ssh_argv(ctx, server, command, stdin=True)) + result = _run_ssh(_ssh_argv(ctx, server, command)) + _write_stdout(result.stdout.decode("utf-8", errors="replace")) + if result.rc == EXIT_NOT_FOUND: + logger.error("%s にジョブ %r のログが無い", server.name, name) + elif result.rc != 0: + _report_remote_failure(result, server) + return _as_exit_code(result.rc) + + +def _cmd_kill(args: argparse.Namespace, ctx: _Context) -> int: + """kill サブコマンド: ジョブを止める（ジョブのあるサーバを自動で探す）.""" + name = validate_job_name(args.name) + probes = _probe_servers(ctx, _target_servers(ctx, args.server)) + _warn_unreachable(probes) + found = servers_with_session(probes, name, jobs_only=True) + if not found and servers_with_session(probes, name, jobs_only=False): + raise ValueError( + f"{name!r} は tias-remote のジョブではない tmux セッションである（止めない）" + ) + unreachable = sum(1 for probe in probes if not probe.reachable) + server = _pick_single( + [probe.server for probe in found], what=f"実行中のジョブ {name!r}", unreachable=unreachable + ) + result = _run_ssh( + _ssh_argv(ctx, server, build_kill_command(server.remote_dir, name)), + timeout=ctx.query_timeout_sec, + ) + if result.rc == EXIT_NO_SESSION: + logger.error("%s のジョブ %r は実行中ではない（既に終了した可能性）", server.name, name) + return result.rc + if result.rc != 0: + _report_remote_failure(result, server) + return _as_exit_code(result.rc) + _write_stdout(f"ジョブ {name} を終了した（{server.name}）\n") + return 0 + + +# --------------------------------------------------------------------------- +# pull +# --------------------------------------------------------------------------- + + +def _locate_run(ctx: _Context, run_id: str, server_name: str | None) -> tuple[ServerConfig, bytes]: + """RUN_ID のディレクトリがあるサーバを探し，そのファイル一覧を返す（複数・0 台ならエラー）.""" + servers = _target_servers(ctx, server_name) + results = _query_servers(ctx, servers, lambda s: build_pull_list_command(s.remote_dir, run_id)) + found: list[tuple[ServerConfig, bytes]] = [] + unreachable = 0 + for server, result in zip(servers, results, strict=True): + if result.rc == 0: + found.append((server, result.stdout)) + elif result.rc == EXIT_UNSAFE_PATH: + raise ValueError( + f"{server.name} の {REMOTE_RUNS_DIR}/{run_id} がリンクであるため回収しない" + ) + elif result.rc != EXIT_NOT_FOUND: + # 一覧の失敗理由（find の標準エラー）はファイル名を含みうるので表示しない + unreachable += 1 + reason = ( + "応答がタイムアウトした" if result.rc == _TIMEOUT_RC else f"終了コード {result.rc}" + ) + logger.warning("%s で実験結果を探せなかった（%s）", server.name, reason) + server = _pick_single([s for s, _ in found], what=f"実験 {run_id!r}", unreachable=unreachable) + listing = next(out for s, out in found if s.name == server.name) + return server, listing + + +def _receive_tar(argv: list[str], payload: bytes, destination: BinaryIO) -> _SshResult: + """ssh の標準入力へ payload を送り，標準出力（tar）を destination へ書き写す. + + 書き込みと読み出しを別スレッドにする（送る一覧が大きいとき，リモートの tar の出力で + パイプが詰まって互いに待つのを避けるため）．標準エラーは表示しない（リモートの tar の + メッセージはファイル名を含みうるため）． + """ + proc = _popen(argv, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + if proc.stdin is None or proc.stdout is None or proc.stderr is None: + raise RuntimeError("ssh の標準入出力を開けなかった") + stdin, stderr_pipe = proc.stdin, proc.stderr + stderr_chunks: list[bytes] = [] + + def _write_payload() -> None: + # ssh が先に終了していれば書き込みは失敗する．終了コードで失敗を報告するので無視する + with contextlib.suppress(OSError): + stdin.write(payload) + with contextlib.suppress(OSError): + stdin.close() + + writer = threading.Thread(target=_write_payload, daemon=True) + reader = threading.Thread(target=lambda: stderr_chunks.append(stderr_pipe.read()), daemon=True) + writer.start() + reader.start() + copy_stream(proc.stdout, destination) + writer.join() + reader.join() + rc = proc.wait() + return _SshResult(rc, b"", b"".join(stderr_chunks)) + + +def _cmd_pull(args: argparse.Namespace, ctx: _Context) -> int: + """pull サブコマンド: 実験結果を回収ホワイトリストに従って持ち帰る.""" + run_id = validate_run_id(args.run_id) + server, listing = _locate_run(ctx, run_id, args.server) + plan = plan_pull(parse_find_listing(listing), run_id, with_checkpoints=args.with_checkpoints) + if not plan.files: + raise ValueError( + f"{server.name} の実験 {run_id!r} に回収対象のファイルが無い" + f"（回収しなかったファイル {plan.skipped} 件）" + ) + + dest = ctx.project_root / REMOTE_RUNS_DIR + logger.info("取得: %s の %s/%s → %s", server.name, REMOTE_RUNS_DIR, run_id, dest / run_id) + payload = b"".join(path.encode("utf-8", "surrogateescape") + b"\0" for path in plan.files) + argv = _ssh_argv(ctx, server, build_pull_tar_command(server.remote_dir), stdin=True) + start = time.monotonic() + with tempfile.TemporaryFile() as archive: + result = _receive_tar(argv, payload, archive) + if result.rc == _TAR_WARNING_RC: + logger.warning("一部のファイルが読み取り中に変化した（ジョブが実行中の可能性）") + elif result.rc == _SSH_FAILURE_RC: + logger.error( + "%s へ ssh 接続できなかった（%s）", server.name, _last_stderr_line(result.stderr) + ) + return result.rc + elif result.rc != 0: + logger.error("%s で結果の書き出しに失敗した（終了コード %d）", server.name, result.rc) + return result.rc + try: + extracted = extract_pull_archive( + archive, dest, run_id, with_checkpoints=args.with_checkpoints + ) + except tarfile.TarError as exc: + raise ValueError("受け取ったアーカイブが壊れているため，何も展開せず中止した") from exc + + if extracted.files != len(plan.files): + logger.warning( + "回収予定 %d 件のうち %d 件しか届かなかった", len(plan.files), extracted.files + ) + skipped = f"回収しなかったファイル: {plan.skipped} 件" + if plan.skipped_checkpoints: + skipped += ( + f"（うち checkpoints/*.pt {plan.skipped_checkpoints} 件は" + " --with-checkpoints で回収できる）" + ) + _write_stdout( + f"取得完了: {extracted.files} ファイル / {_format_size(extracted.bytes)}" + f"（{time.monotonic() - start:.1f} 秒）→ {REMOTE_RUNS_DIR}/{run_id}/\n{skipped}\n" + ) + return 0 + + +# --------------------------------------------------------------------------- +# 引数 +# --------------------------------------------------------------------------- + + +def _parse_args(argv: Sequence[str] | None) -> argparse.Namespace: + """コマンドライン引数をパースして返す.""" + parser = argparse.ArgumentParser( + prog="tias-remote", + description=( + "研究室の GPU サーバ（configs/remote.yaml）で空いている GPU を選んで実行する" + "（コード送信・ジョブ起動・状態確認・結果回収）．" + ), + epilog=_EPILOG, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + common = argparse.ArgumentParser(add_help=False) + common.add_argument( + "--config", + type=Path, + default=None, + metavar="PATH", + help=f"サーバ設定 YAML（既定: プロジェクトルートの {DEFAULT_CONFIG_PATH.as_posix()}）．", + ) + server_opt = argparse.ArgumentParser(add_help=False) + server_opt.add_argument( + "--server", default=None, metavar="NAME", help="対象のサーバ（既定: 全サーバから探す）．" + ) + sub = parser.add_subparsers(dest="subcommand", required=True, metavar="SUBCOMMAND") + + p_gpus = sub.add_parser( + "gpus", + parents=[common], + help="全サーバの GPU 一覧（状態・使い方の型・利用者）と run が選ぶ GPU を表示する", + description=( + "各 GPU を一定時間観測し，観測中の最大値で 未使用（閾値内・計算プロセス無し）／" + "空き（利用者あり）／使用中／状態不明 に分けて表示する．run は全候補サーバを通して" + "まず未使用を，無ければ空き（利用者あり）を選ぶ．" + ), + ) + p_gpus.add_argument( + "--observe", + type=int, + default=None, + metavar="SEC", + help="観測秒数（設定の observe_sec を上書き．0 で瞬間値のみ．0〜120）．", + ) + p_gpus.add_argument( + "--no-data", + action="store_true", + help="データ未配置（data_ready: false）のサーバも run の候補として選ぶ．", + ) + + sub.add_parser("status", parents=[common], help="全サーバの実行中ジョブと GPU の状況を表示する") + + p_push = sub.add_parser( + "push", + parents=[common, server_opt], + help="コードをサーバへ送る（省略時は接続できる全サーバ）", + description=( + "リポジトリを tar over ssh でサーバの remote_dir へ上書き展開する．" + ".git / .venv / キャッシュ / data / runs / .claude は送らない．" + ), + ) + p_push.add_argument( + "--dry-run", action="store_true", help="送信せず，対象ファイルの一覧だけ表示する．" + ) + + p_run = sub.add_parser( + "run", + parents=[common, server_opt], + help="GPU を選び（未使用を最優先），コードを送って tmux でジョブを起動する", + description=( + "GPU を選び（data_ready のサーバ全体を通して，まず未使用を，無ければ" + "空き（利用者あり）を priority 順・GPU 番号順に探す），" + "そのサーバへコードを送ってから remote_dir で `uv run ` を tmux で起動する．" + f"出力は remote_dir/{REMOTE_LOG_DIR}/.log に追記され，接続が切れても継続する．" + "コマンドは -- の後に書く．" + ), + ) + p_run.add_argument( + "--name", + required=True, + metavar="NAME", + help="ジョブ名（tmux セッション名・ログ名）．英数字で始まり英数字・_・-，64 文字以内．", + ) + p_run.add_argument( + "--gpu", + type=int, + default=None, + metavar="I", + help="使う GPU 番号（--server が必須．使用中・状態不明ならエラー）．", + ) + p_run.add_argument( + "--force", + action="store_true", + help="--gpu で指定した GPU が使用中・状態不明でも使う．", + ) + p_run.add_argument( + "--no-data", + action="store_true", + help="データ未配置（data_ready: false）のサーバも候補にする．", + ) + p_run.add_argument("--no-push", action="store_true", help="起動前にコードを送らない．") + p_run.add_argument( + "--observe", + type=int, + default=None, + metavar="SEC", + help="GPU を選ぶ前の観測秒数（設定の observe_sec を上書き．0 で瞬間値のみ）．", + ) + p_run.add_argument( + "--env", + action="append", + default=[], + metavar="KEY=VALUE", + help="ジョブに追加で渡す環境変数．複数指定可．", + ) + p_run.add_argument( + "--dry-run", + action="store_true", + help="GPU を選ぶところまで行い，送信・起動せずにリモートコマンドを表示する．", + ) + p_run.add_argument( + "command", + nargs=argparse.REMAINDER, + metavar="-- COMMAND", + help="サーバで実行するコマンド（uv run の後ろに付く部分）．", + ) + + p_logs = sub.add_parser( + "logs", parents=[common, server_opt], help="ジョブのログを表示する（サーバを自動で探す）" + ) + p_logs.add_argument("name", metavar="NAME", help="ジョブ名．") + p_logs.add_argument( + "--lines", "-n", type=int, default=40, metavar="N", help="表示する末尾の行数（既定 40）．" + ) + p_logs.add_argument( + "--follow", "-f", action="store_true", help="追記を追従表示する（Ctrl+C で終了）．" + ) + + p_kill = sub.add_parser( + "kill", parents=[common, server_opt], help="ジョブを止める（サーバを自動で探す）" + ) + p_kill.add_argument("name", metavar="NAME", help="ジョブ名．") + + p_pull = sub.add_parser( + "pull", + parents=[common, server_opt], + help="実験結果を回収ホワイトリストに従って持ち帰る", + description=( + "サーバの runs// のうち回収ホワイトリスト（config.yaml・history.csv・" + "metrics.json・confusion_matrix.csv・log.txt・figures/*.png）に一致するファイルだけを" + "手元の runs// へ持ち帰る（同名ファイルは上書き）．" + ), + ) + p_pull.add_argument("run_id", metavar="RUN_ID", help="実験 ID（runs/ 直下のディレクトリ名）．") + p_pull.add_argument( + "--with-checkpoints", action="store_true", help="checkpoints/*.pt も持ち帰る．" + ) + + return parser.parse_args(argv) + + +_HANDLERS: dict[str, Callable[[argparse.Namespace, _Context], int]] = { + "gpus": _cmd_gpus, + "status": _cmd_status, + "push": _cmd_push, + "run": _cmd_run, + "logs": _cmd_logs, + "kill": _cmd_kill, + "pull": _cmd_pull, +} + + +def _build_context(args: argparse.Namespace) -> _Context: + """設定を読み込み，実行環境を組み立てる.""" + project_root = find_project_root(Path.cwd()) + config_path = args.config if args.config is not None else project_root / DEFAULT_CONFIG_PATH + config = load_remote_config(config_path) + return _Context( + config=config, + ssh=resolve_ssh_executable(os.environ), + jump_host=resolve_jump_host(config, os.environ), + project_root=project_root, + ) + + +def main(argv: Sequence[str] | None = None) -> int: + """tias-remote コマンドのメインエントリポイント. + + Args: + argv: 引数リスト．None のとき sys.argv[1:] を使用． + + Returns: + 終了コード（成功時 0．ssh・リモートの失敗時はその終了コード， + 入力エラー・空き GPU 無しは 1）． + """ + args = _parse_args(argv) + _configure_logging() + try: + ctx = _build_context(args) + if not is_self_check_disabled(os.environ): + self_server = detect_self_server(ctx.config, socket.gethostname()) + if self_server is not None: + logger.error( + "ここはサーバ %s である（ホスト名が一致）．" + "tias-remote は手元の PC から実行する．" + "サーバ上で直接回すなら tmux の中で uv run ... を実行する．" + "判定が誤りなら環境変数 %s=1 で無効にできる", + self_server.name, + ALLOW_SELF_ENV_VAR, + ) + return 1 + if ctx.jump_host is None and JUMP_ENV_VAR in os.environ: + logger.info("%s が空のため，踏み台を使わず直接接続する", JUMP_ENV_VAR) + return _HANDLERS[args.subcommand](args, ctx) + except ValidationError as exc: + logger.error("サーバ設定が不正である:\n%s", exc) + return 1 + except ValueError as exc: + logger.error("%s", exc) + return 1 + except KeyboardInterrupt: + logger.error("中断した") + return 130 + + +if __name__ == "__main__": + # uv の生成するランチャ（tias-remote.exe）が Windows の Smart App Control に + # ブロックされた場合の代替として `uv run python -m tias_judge_model.cli.remote` を使える + sys.exit(main()) diff --git a/src/tias_judge_model/remote/__init__.py b/src/tias_judge_model/remote/__init__.py new file mode 100644 index 0000000..49602ad --- /dev/null +++ b/src/tias_judge_model/remote/__init__.py @@ -0,0 +1,5 @@ +"""リモート実行: 研究室の GPU サーバへのコード送信・空き GPU の選択・ジョブ実行・結果回収. + +ssh を呼ぶ副作用は `tias_judge_model.cli.remote` に置き，このパッケージは設定の検証・ +コマンド文字列の組み立て・出力の解析・選択・回収ホワイトリストの判定などを担う（SPEC_03）． +""" diff --git a/src/tias_judge_model/remote/commands.py b/src/tias_judge_model/remote/commands.py new file mode 100644 index 0000000..cb81044 --- /dev/null +++ b/src/tias_judge_model/remote/commands.py @@ -0,0 +1,598 @@ +"""サーバ上で実行するシェルコマンド文字列と ssh 引数を組み立てる純関数群. + +`ssh @
''` で送る文字列を作る．リモートのログインシェル（Ubuntu では +bash）が解釈する前提で POSIX sh の範囲で書き，ジョブ本体だけは `${PIPESTATUS[0]}` を使うため +`bash -c` で包む．ユーザー由来の値（パス・コマンド引数）はすべて `shlex.quote` でクォートし， +ジョブ名・環境変数名・RUN_ID は書式を検証してから埋め込む． + +Windows から ssh へ渡す引数に改行を含めないため，複数の文は `; ` で 1 行につなぐ． +リモート側が出すメッセージは ASCII に限る（ssh の stderr は Windows コンソールへ素通しされ， +コードページによっては日本語が化けるため．日本語の説明は呼び出し側の CLI が出す）． +""" + +from __future__ import annotations + +import re +import shlex +from collections.abc import Mapping, Sequence + +from tias_judge_model.schema import DATA_ROOT_ENV_VAR + +REMOTE_RUNS_DIR = "runs" +REMOTE_LOG_DIR = f"{REMOTE_RUNS_DIR}/_remote_logs" + +# リモートスクリプトが返す終了コード．ssh 自身の 255・tar の 1/2 と衝突しない値にする +EXIT_REMOTE_DIR_MISSING = 90 +EXIT_SESSION_EXISTS = 91 +EXIT_NOT_FOUND = 92 +EXIT_NO_SESSION = 93 +EXIT_UNSAFE_PATH = 94 + +# probe コマンドの出力を区切る目印（ログインシェルが余計な出力をしても区切りで読み分ける） +PROBE_SESSIONS_MARKER = "@@tias-sessions" +PROBE_GPUS_MARKER = "@@tias-gpus" +PROBE_GPU_RC_PREFIX = "@@tias-gpu-rc " +PROBE_SAMPLES_MARKER = "@@tias-samples" +PROBE_PMON_MARKER = "@@tias-pmon" +PROBE_APPS_MARKER = "@@tias-apps" +PROBE_APPS_RC_PREFIX = "@@tias-apps-rc " +PROBE_OWNERS_MARKER = "@@tias-owners" +PROBE_END_MARKER = "@@tias-end" +PROBE_JOB_PREFIX = "job " +PROBE_OTHER_PREFIX = "other " + +NVIDIA_SMI_QUERY = ( + "nvidia-smi --query-gpu=index,uuid,name,memory.used,memory.total,utilization.gpu" + " --format=csv,noheader,nounits" +) +# 観測用のサンプリング（-lms <間隔> を付けて繰り返し出力させ，timeout で打ち切る） +NVIDIA_SMI_SAMPLE_QUERY = ( + "nvidia-smi --query-gpu=index,utilization.gpu,memory.used --format=csv,noheader,nounits" +) +# プロセス別の SM 使用率（1 秒間隔．-c <回数> を付けて観測秒数ぶん取る） +NVIDIA_SMI_PMON_COMMAND = "nvidia-smi pmon -s u -d 1" +# 計算プロセスのみ（画面描画の Xorg 等は含まない）．gpu_uuid で上の GPU と対応付ける +NVIDIA_SMI_APPS_QUERY = ( + "nvidia-smi --query-compute-apps=gpu_uuid,pid,used_memory --format=csv,noheader,nounits" +) + +# tmux のセッション名に使えない "." ":" を避け，ログのファイル名にもそのまま使える文字に限る +_JOB_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_-]{0,63}$") +_ENV_KEY_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") +# ssh の接続先（ユーザー名・アドレス・踏み台）に許す文字．空白や `-` 始まりでオプションとして +# 解釈されるのを防ぐ +_SSH_TOKEN_RE = re.compile(r"^[A-Za-z0-9_.][A-Za-z0-9_.@:%-]*$") + +_NOW = "\"$(date '+%Y-%m-%d %H:%M:%S')\"" + + +def quote_remote_path(path: str) -> str: + """リモートのパスをシェル用にクォートする. + + 先頭の `~` はリモートの `$HOME` に展開させる必要があるため，`~/...` は + `"$HOME"/'...'` の形に変換する（丸ごとクォートするとチルダ展開が起きない）． + + Args: + path: リモートのパス（`~/...`・絶対パス・ホーム基準の相対パス）． + + Returns: + シェルにそのまま埋め込める文字列． + + Raises: + ValueError: 空文字，または `~ユーザー名/...` 形式の場合． + """ + if not path or not path.strip(): + raise ValueError("リモートパスが空である") + if path == "~": + return '"$HOME"' + if path.startswith("~/"): + rest = path[2:] + return f'"$HOME"/{shlex.quote(rest)}' if rest else '"$HOME"' + if path.startswith("~"): + raise ValueError( + f"~ユーザー名 形式のリモートパスには対応していない: {path!r}（~/... か絶対パスを指定）" + ) + return shlex.quote(path) + + +def validate_ssh_token(value: str, *, what: str) -> str: + """ssh の接続先に使う文字列（ユーザー名・アドレス・踏み台）を検証して返す. + + Args: + value: 検証する文字列． + what: エラーメッセージに出す項目名． + + Returns: + 検証済みの文字列． + + Raises: + ValueError: 空・`-` 始まり・空白や記号を含む場合． + """ + if not _SSH_TOKEN_RE.match(value): + raise ValueError(f"{what} が不正である: {value!r}（英数字・. _ - @ : % のみ）") + return value + + +def validate_job_name(name: str) -> str: + """ジョブ名（tmux セッション名・ログのファイル名）を検証して返す. + + Args: + name: ジョブ名． + + Returns: + 検証済みのジョブ名． + + Raises: + ValueError: 英数字で始まり英数字・`_`・`-` のみ（1〜64 文字）でない場合． + """ + if not _JOB_NAME_RE.match(name): + raise ValueError( + f"ジョブ名が不正である: {name!r}（英数字で始まり，英数字・_・- のみ，64 文字以内）" + ) + return name + + +def validate_env_key(key: str) -> str: + """環境変数名を検証して返す. + + Args: + key: 環境変数名． + + Returns: + 検証済みの環境変数名． + + Raises: + ValueError: 環境変数名として不正な場合． + """ + if not _ENV_KEY_RE.match(key): + raise ValueError(f"環境変数名が不正である: {key!r}") + return key + + +def parse_env_assignment(text: str) -> tuple[str, str]: + """`KEY=VALUE` 形式の環境変数指定を分解する. + + Args: + text: `KEY=VALUE` 形式の文字列（VALUE は空でもよい）． + + Returns: + (KEY, VALUE)． + + Raises: + ValueError: `=` が無い，または KEY が環境変数名として不正な場合． + """ + key, sep, value = text.partition("=") + if not sep or not _ENV_KEY_RE.match(key): + raise ValueError(f"環境変数の指定が不正である: {text!r}（KEY=VALUE 形式）") + return key, value + + +def build_job_env(gpu_index: int, extra: Sequence[str]) -> dict[str, str]: + """ジョブに渡す環境変数（値をそのまま渡すもの）を組み立てる. + + `CUDA_DEVICE_ORDER=PCI_BUS_ID` を付けるのは，CUDA の既定の番号付け（速い順）が + nvidia-smi の index（PCI バス順）と異なりうるため（異なると空きと判定した GPU とは別の + GPU を使ってしまう）．`TIAS_DATA_ROOT` はパスとして別に渡す（`build_job_script` の + `path_env`）．GPU の割り当てを迂回させないため，`--env` で予約済みの変数は上書きできない． + + Args: + gpu_index: 選んだ GPU の番号（nvidia-smi の index）． + extra: `--env` で指定された `KEY=VALUE`（後勝ち）． + + Returns: + 環境変数の辞書（挿入順＝ログに表示する順）． + + Raises: + ValueError: GPU 番号が負，`KEY=VALUE` が不正，または予約済みの変数を指定した場合． + """ + if gpu_index < 0: + raise ValueError(f"GPU 番号は 0 以上: {gpu_index}") + env = { + "CUDA_DEVICE_ORDER": "PCI_BUS_ID", + "CUDA_VISIBLE_DEVICES": str(gpu_index), + # パイプ越しでも Python の出力がすぐログに出るようにする（既定はブロックバッファ） + "PYTHONUNBUFFERED": "1", + } + reserved = {*env, DATA_ROOT_ENV_VAR} + for item in extra: + key, value = parse_env_assignment(item) + if key in reserved: + raise ValueError(f"{key} は tias-remote が設定するため --env では指定できない") + env[key] = value + return env + + +def build_ssh_argv( + ssh: str, + *, + user: str, + address: str, + jump_host: str | None, + connect_timeout_sec: int, + remote_command: str, + stdin: bool = False, +) -> list[str]: + """ssh を起動する引数リストを返す. + + 常に `BatchMode=yes`（パスワード入力を求めず失敗させる）と `ConnectTimeout` を付ける． + 標準入力を使わないコマンドには `-n` を付ける（標準入力が開いたままだと応答せず止まる + ことがあるため．ENV_02 トラブルシューティング）． + + Args: + ssh: ssh 実行ファイル． + user: 接続ユーザー． + address: 接続先アドレス． + jump_host: 踏み台（`-J` に渡す）．None なら直接接続． + connect_timeout_sec: 接続タイムアウト（秒）． + remote_command: リモートで実行するシェルコマンド（1 引数として渡す）． + stdin: True なら標準入力をリモートへ渡す（`-n` を付けない）． + + Returns: + subprocess に渡す引数リスト． + + Raises: + ValueError: 接続先の書式が不正，またはタイムアウトが 1 未満の場合． + """ + if connect_timeout_sec < 1: + raise ValueError(f"接続タイムアウトは 1 以上でなければならない（{connect_timeout_sec}）") + destination = ( + f"{validate_ssh_token(user, what='ユーザー名')}@" + f"{validate_ssh_token(address, what='アドレス')}" + ) + argv = [ssh, "-o", "BatchMode=yes", "-o", f"ConnectTimeout={connect_timeout_sec}"] + if jump_host is not None: + argv += ["-J", validate_ssh_token(jump_host, what="踏み台")] + if not stdin: + argv.append("-n") + argv += [destination, remote_command] + return argv + + +def _cd_command(remote_dir: str) -> str: + """リモートディレクトリへ移動し，無ければ専用の終了コードで抜ける文を返す.""" + quoted = quote_remote_path(remote_dir) + return ( + f"cd {quoted} || {{ echo 'tias-remote: remote directory not found:' {quoted} >&2; " + f"exit {EXIT_REMOTE_DIR_MISSING}; }}" + ) + + +def job_log_path(name: str) -> str: + """ジョブログのパス（remote_dir 基準．クォート前）を返す. + + Args: + name: 検証済みのジョブ名（`validate_job_name`）． + + Returns: + `runs/_remote_logs/.log`． + """ + return f"{REMOTE_LOG_DIR}/{name}.log" + + +def _tmux_target(name: str) -> str: + """tmux の `-t` に渡すセッション指定をクォートして返す．name は検証済みとする. + + `=` を前置して完全一致にする（付けないと tmux は前方一致で別のジョブを選びうる）． + """ + return shlex.quote(f"={name}") + + +def build_probe_command(remote_dir: str, *, observe_sec: int, sample_interval_ms: int) -> str: + """tmux セッション・GPU の観測・計算プロセスとその所有者を 1 回の ssh で取るコマンドを返す. + + 出力の形式（`parse_probe_output` が解釈する）: + + - `@@tias-sessions` の後に，tmux セッション 1 件につき 1 行．このツールのジョブログ + （`remote_dir/runs/_remote_logs/<名前>.log`）があるものは `job <名前>`，無いものは + `other <名前>`（同じユーザーの別用途のセッション） + - `@@tias-gpus` の後に GPU の一覧（`nvidia-smi --query-gpu` の CSV），続けて + `@@tias-gpu-rc <終了コード>`（nvidia-smi が無ければ 127 で，以降の節は出ない） + - `observe_sec > 0` のとき，`@@tias-samples` の後に観測中のサンプル（`index, 使用率, + 使用中メモリ` の CSV．`sample_interval_ms` 間隔），`@@tias-pmon` の後に + `nvidia-smi pmon -s u` の出力（1 秒間隔．失敗・未対応なら空） + - `@@tias-apps` の後に計算プロセス（`nvidia-smi --query-compute-apps` の CSV），続けて + `@@tias-apps-rc <終了コード>` + - `@@tias-owners` の後に，計算プロセス 1 件につき ` <所有ユーザー> <経過秒>` + （ps で取れなければ pid だけ） + - 最後に `@@tias-end`（出力が途中で切れていないことの確認用） + + remote_dir が無くても失敗しない（まだ push していないサーバも GPU 一覧に出すため）． + 所要時間はおよそ observe_sec 秒（観測はサーバ側で行う）． + + Args: + remote_dir: リモートのプロジェクトルート． + observe_sec: 観測秒数（0 なら観測せず瞬間値だけを取る）． + sample_interval_ms: 観測の間隔（ミリ秒）． + + Returns: + リモートで実行するシェルコマンド． + + Raises: + ValueError: 観測秒数が負，または間隔が 1 未満の場合． + """ + log_dir = f"{quote_remote_path(remote_dir)}/{REMOTE_LOG_DIR}" + sessions = ( + "tmux ls -F '#{session_name}' 2>/dev/null | while IFS= read -r s; do " + f'if [ -f {log_dir}/"$s".log ]; then echo "{PROBE_JOB_PREFIX}$s"; ' + f'else echo "{PROBE_OTHER_PREFIX}$s"; fi; done' + ) + if observe_sec < 0 or sample_interval_ms < 1: + raise ValueError( + f"観測の設定が不正である（observe_sec={observe_sec}, " + f"sample_interval_ms={sample_interval_ms}）" + ) + # 所有者は ps の user 列を広げて取る（既定幅では長いユーザー名が UID 表記・切り詰めになる）． + # etimes は起動からの経過秒数（他ユーザーのプロセスでも取れる） + owners = ( + "printf '%s\\n' \"$apps\" | while IFS=, read -r _u p _m; do " + "p=$(printf '%s' \"$p\" | tr -d ' '); " + "case \"$p\" in ''|*[!0-9]*) continue;; esac; " + 'echo "$p $(ps -o user:64=,etimes= -p "$p" 2>/dev/null | tr -s \' \')"; done' + ) + observe = "" + if observe_sec > 0: + # GPU ごとのサンプリングと pmon（プロセス別 SM 使用率）を並行して走らせる．出力が + # 混ざらないよう一時ディレクトリのファイルに分けて書き，終わってから順に出す． + # timeout で打ち切った nvidia-smi -lms は 124 で終わるが，出力は有効なので終了コードは + # 見ない（サンプルが 0 件の GPU を状態不明にする）．pmon が止まらない場合に備えて + # pmon にも余裕を持たせた timeout を付ける + observe = ( + 't=$(mktemp -d 2>/dev/null); if [ -n "$t" ]; then ' + f"timeout {observe_sec + 5} {NVIDIA_SMI_PMON_COMMAND} -c {observe_sec}" + ' > "$t/pmon" 2>/dev/null & pp=$!; ' + f"timeout {observe_sec} {NVIDIA_SMI_SAMPLE_QUERY} -lms {sample_interval_ms}" + ' > "$t/samples" 2>/dev/null; wait "$pp"; ' + f"echo '{PROBE_SAMPLES_MARKER}'; cat \"$t/samples\"; " + f"echo '{PROBE_PMON_MARKER}'; cat \"$t/pmon\"; " + 'rm -rf "$t"; fi; ' + ) + gpus = ( + "if command -v nvidia-smi >/dev/null 2>&1; then " + f'{NVIDIA_SMI_QUERY} 2>/dev/null; echo "{PROBE_GPU_RC_PREFIX}$?"; ' + f"{observe}" + f"apps=$({NVIDIA_SMI_APPS_QUERY} 2>/dev/null); arc=$?; " + f"echo '{PROBE_APPS_MARKER}'; if [ -n \"$apps\" ]; then printf '%s\\n' \"$apps\"; fi; " + f'echo "{PROBE_APPS_RC_PREFIX}$arc"; ' + f"echo '{PROBE_OWNERS_MARKER}'; {owners}; " + f"else echo '{PROBE_GPU_RC_PREFIX}127'; fi" + ) + return "; ".join( + [ + f"echo '{PROBE_SESSIONS_MARKER}'", + sessions, + f"echo '{PROBE_GPUS_MARKER}'", + gpus, + f"echo '{PROBE_END_MARKER}'", + ] + ) + + +def build_push_command(remote_dir: str) -> str: + """標準入力の tar ストリームをリモートディレクトリへ展開するコマンドを返す. + + Args: + remote_dir: 展開先（リモートのプロジェクトルート）． + + Returns: + リモートで実行するシェルコマンド． + """ + quoted = quote_remote_path(remote_dir) + return f"mkdir -p {quoted} && tar -xf - -C {quoted}" + + +def build_job_script( + remote_dir: str, + name: str, + command: Sequence[str], + *, + uv_path: str, + server_name: str, + gpu_index: int, + env: Mapping[str, str], + path_env: Mapping[str, str], +) -> str: + """tmux セッション内で bash が実行するジョブ本体のスクリプトを返す. + + 出力は `runs/_remote_logs/.log` に追記し（前回の同名ジョブのログも残す）， + 開始行（サーバ名・GPU・日時）と終了コード行で挟む．標準入力は /dev/null にする + （入力待ちで止まらないように）．非対話 ssh では `~/.local/bin` が PATH に無いため， + uv のディレクトリと `$HOME/.local/bin` を PATH に足す． + + Args: + remote_dir: リモートのプロジェクトルート． + name: ジョブ名． + command: `uv run` の後ろに付けて実行するコマンドと引数． + uv_path: リモートの uv 実行ファイル． + server_name: ログの開始行に書くサーバ名． + gpu_index: ログの開始行に書く GPU 番号． + env: ジョブに渡す環境変数（値はそのまま渡す）． + path_env: ジョブに渡す環境変数のうち，値がリモートのパス（`~` を `$HOME` に展開する） + であるもの． + + Returns: + bash スクリプト（1 行）． + + Raises: + ValueError: コマンドが空，またはジョブ名・環境変数名・パスが不正な場合． + """ + if not command: + raise ValueError("実行するコマンドが空である") + log = job_log_path(validate_job_name(name)) + for key in [*env, *path_env]: + validate_env_key(key) + + uv_quoted = quote_remote_path(uv_path) + run_argv = [uv_quoted, "run", *(shlex.quote(arg) for arg in command)] + display = shlex.join( + [f"{k}={v}" for k, v in [*env.items(), *path_env.items()]] + ["uv", "run", *command] + ) + header = f"server={server_name} gpu={gpu_index}" + statements = [ + _cd_command(remote_dir), + f"log={shlex.quote(log)}", + f"printf '=== [%s] start %s host=%s: %s\\n' {_NOW} {shlex.quote(header)} " + f'"$(hostname)" {shlex.quote(display)} >> "$log"', + f'export PATH="$(dirname -- {uv_quoted})":"$HOME/.local/bin:$PATH"', + *[f"export {k}={shlex.quote(v)}" for k, v in env.items()], + *[f"export {k}={quote_remote_path(v)}" for k, v in path_env.items()], + f'{" ".join(run_argv)} < /dev/null 2>&1 | tee -a "$log"', + "rc=${PIPESTATUS[0]}", + f'printf \'=== [%s] exit %s\\n\' {_NOW} "$rc" >> "$log"', + 'exit "$rc"', + ] + return "; ".join(statements) + + +def build_run_command(remote_dir: str, name: str, job_script: str) -> str: + """ジョブ本体を tmux の新規セッション（デタッチ状態）で起動するコマンドを返す. + + 同名セッションが既にあれば起動せず `EXIT_SESSION_EXISTS` で終わる． + tmux のサーバはデーモン化して ssh の入出力を閉じるため，ssh はすぐ戻り， + 接続が切れてもジョブは継続する． + + Args: + remote_dir: リモートのプロジェクトルート． + name: ジョブ名（tmux セッション名）． + job_script: `build_job_script` が返したスクリプト． + + Returns: + リモートで実行するシェルコマンド． + + Raises: + ValueError: ジョブ名が不正な場合． + """ + name = validate_job_name(name) + inner = "bash -c " + shlex.quote(job_script) + statements = [ + _cd_command(remote_dir), + f"mkdir -p {REMOTE_LOG_DIR} || exit 1", + f"if tmux has-session -t {_tmux_target(name)} 2>/dev/null; then " + f"echo 'tias-remote: tmux session already exists:' {shlex.quote(name)} >&2; " + f"exit {EXIT_SESSION_EXISTS}; fi", + f'tmux new-session -d -s {shlex.quote(name)} -c "$PWD" {shlex.quote(inner)}', + ] + return "; ".join(statements) + + +def build_log_exists_command(remote_dir: str, name: str) -> str: + """ジョブログがあれば 0，無ければ `EXIT_NOT_FOUND` で終わるコマンドを返す. + + Args: + remote_dir: リモートのプロジェクトルート． + name: ジョブ名． + + Returns: + リモートで実行するシェルコマンド． + + Raises: + ValueError: ジョブ名が不正な場合． + """ + log = shlex.quote(job_log_path(validate_job_name(name))) + return ( + f"cd {quote_remote_path(remote_dir)} 2>/dev/null && [ -f {log} ] || exit {EXIT_NOT_FOUND}" + ) + + +def build_logs_command(remote_dir: str, name: str, *, lines: int, follow: bool) -> str: + """ジョブログの末尾を表示する（follow なら追従する）コマンドを返す. + + 追従モードでは，ssh の標準入力が閉じたら（手元で Ctrl+C・終了したら）tail を止める． + 疑似端末なしの ssh では切断しても tail に SIGHUP が届かず，リモートに残り続けるため． + + Args: + remote_dir: リモートのプロジェクトルート． + name: ジョブ名． + lines: 表示する末尾の行数（0 以上）． + follow: True なら追記を追従する． + + Returns: + リモートで実行するシェルコマンド． + + Raises: + ValueError: ジョブ名が不正，または lines が負の場合． + """ + if lines < 0: + raise ValueError(f"表示行数は 0 以上でなければならない（{lines}）") + log = shlex.quote(job_log_path(validate_job_name(name))) + check = ( + f"[ -f {log} ] || {{ echo 'tias-remote: log not found:' {log} >&2; " + f"exit {EXIT_NOT_FOUND}; }}" + ) + if follow: + tail = f'tail -n {lines} -F {log} & t=$!; cat >/dev/null; kill "$t" 2>/dev/null' + else: + tail = f"tail -n {lines} {log}" + return "; ".join([_cd_command(remote_dir), check, tail]) + + +def build_kill_command(remote_dir: str, name: str) -> str: + """ジョブの tmux セッションを終了し，ログにその旨を追記するコマンドを返す. + + tmux がセッションを閉じるとジョブの bash は SIGHUP で終わり，自分では終了行を + 書けないため，終了前にここで追記する． + + Args: + remote_dir: リモートのプロジェクトルート． + name: ジョブ名． + + Returns: + リモートで実行するシェルコマンド． + + Raises: + ValueError: ジョブ名が不正な場合． + """ + name = validate_job_name(name) + target = _tmux_target(name) + log = shlex.quote(job_log_path(name)) + statements = [ + f"tmux has-session -t {target} 2>/dev/null || " + f"{{ echo 'tias-remote: no such tmux session:' {shlex.quote(name)} >&2; " + f"exit {EXIT_NO_SESSION}; }}", + f"if cd {quote_remote_path(remote_dir)} 2>/dev/null && [ -f {log} ]; then " + f"printf '=== [%s] killed by tias-remote kill\\n' {_NOW} >> {log}; fi", + f"tmux kill-session -t {target}", + ] + return "; ".join(statements) + + +def build_pull_list_command(remote_dir: str, run_id: str) -> str: + """`runs//` 配下のディレクトリ以外のエントリを種類付きで列挙するコマンドを返す. + + 出力は `<種類> ` を NUL で区切ったもの（GNU find の `-printf '%y %p\\0'`． + 種類は `f` 通常ファイル・`l` シンボリックリンク 等）．回収対象の絞り込みは手元の + `plan_pull` が行い，ここでは何も送らない（一覧はファイル名を含むため表示しない）． + RUN_ID のディレクトリが無ければ `EXIT_NOT_FOUND`，リンクなら `EXIT_UNSAFE_PATH` で終わる． + + Args: + remote_dir: リモートのプロジェクトルート． + run_id: `validate_run_id` で検証済みの RUN_ID． + + Returns: + リモートで実行するシェルコマンド． + """ + run = shlex.quote(run_id) + statements = [ + f"cd {quote_remote_path(remote_dir)}/{REMOTE_RUNS_DIR} 2>/dev/null" + f" || exit {EXIT_NOT_FOUND}", + f"if [ -L {run} ]; then exit {EXIT_UNSAFE_PATH}; fi", + f"[ -d {run} ] || exit {EXIT_NOT_FOUND}", + f"find {run} -mindepth 1 ! -type d -printf '%y %p\\0'", + ] + return "; ".join(statements) + + +def build_pull_tar_command(remote_dir: str) -> str: + """標準入力の NUL 区切りのパス（`runs/` 基準）だけを tar ストリームで書き出すコマンドを返す. + + `--no-recursion` で指定したパス以外を含めない（ディレクトリを渡されても中身を入れない）． + + Args: + remote_dir: リモートのプロジェクトルート． + + Returns: + リモートで実行するシェルコマンド． + """ + return ( + f"cd {quote_remote_path(remote_dir)}/{REMOTE_RUNS_DIR} || exit {EXIT_NOT_FOUND}; " + "tar --null --no-recursion -T - -cf -" + ) diff --git a/src/tias_judge_model/remote/config.py b/src/tias_judge_model/remote/config.py new file mode 100644 index 0000000..8556cb6 --- /dev/null +++ b/src/tias_judge_model/remote/config.py @@ -0,0 +1,274 @@ +"""リモート実行のサーバ設定（`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 diff --git a/src/tias_judge_model/remote/files.py b/src/tias_judge_model/remote/files.py new file mode 100644 index 0000000..981bf00 --- /dev/null +++ b/src/tias_judge_model/remote/files.py @@ -0,0 +1,99 @@ +"""push で送るファイルの選別と，手元のプロジェクトルートの探索. + +リポジトリ内のコード一式を送り，キャッシュ・仮想環境・データ・実験結果・Claude の設定は +送らない（SPEC_03「push」）． +""" + +from __future__ import annotations + +import os +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path, PurePosixPath + +# どの階層にあっても送らないディレクトリ名（.venv は Windows 用でサーバでは使えない） +EXCLUDED_DIR_NAMES = frozenset({".git", ".venv", "__pycache__", ".pytest_cache", ".ruff_cache"}) +# プロジェクト直下でだけ除外するエントリ．data/ は臨床データ・合成データの置き場，runs/ は +# サーバ側の実験結果を上書きしないため，.claude/ は手元の Claude Code 設定のため送らない． +# 直下に限るのは src/tias_judge_model/data/（パッケージ）・tests/data/ を送るため +EXCLUDED_ROOT_ENTRIES = frozenset({"data", "runs", ".claude"}) + +_PACKAGE_DIR = Path("src") / "tias_judge_model" + + +@dataclass(frozen=True) +class PushEntry: + """送信するファイル 1 件. + + Attributes: + path: 手元の実ファイル． + arcname: リモートのプロジェクトルート基準の相対パス（POSIX 区切り）． + size: バイト数． + """ + + path: Path + arcname: str + size: int + + +def find_project_root(start: Path) -> Path: + """start から上方向にたどり，このリポジトリのルートを探す. + + Args: + start: 探索を始めるディレクトリ（通常はカレントディレクトリ）． + + Returns: + `pyproject.toml` と `src/tias_judge_model/` を持つディレクトリ． + + Raises: + ValueError: 見つからない場合． + """ + start = start.resolve() + for directory in (start, *start.parents): + if (directory / "pyproject.toml").is_file() and (directory / _PACKAGE_DIR).is_dir(): + return directory + raise ValueError( + f"{start} から上にプロジェクトルート（pyproject.toml と src/tias_judge_model/）が" + "見つからない．リポジトリ内で実行する" + ) + + +def is_excluded(rel_path: str) -> bool: + """プロジェクトルート基準の相対パスが push から除外されるかを返す. + + Args: + rel_path: POSIX 区切りの相対パス（ファイルでもディレクトリでもよい）． + + Returns: + 除外対象なら True． + """ + parts = PurePosixPath(rel_path).parts + if not parts: + return False + return parts[0] in EXCLUDED_ROOT_ENTRIES or any(part in EXCLUDED_DIR_NAMES for part in parts) + + +def _walk_files(root: Path) -> Iterator[PushEntry]: + """root 配下のファイルを，除外規則を適用しつつ決定的な順で列挙する.""" + for dirpath, dirnames, filenames in os.walk(root): + dir_rel = Path(dirpath).relative_to(root).as_posix() + prefix = "" if dir_rel == "." else f"{dir_rel}/" + dirnames[:] = sorted(d for d in dirnames if not is_excluded(prefix + d)) + for filename in sorted(filenames): + rel = prefix + filename + path = Path(dirpath) / filename + if is_excluded(rel) or not path.is_file(): + continue + yield PushEntry(path=path, arcname=rel, size=path.stat().st_size) + + +def collect_push_entries(root: Path) -> list[PushEntry]: + """送信するファイルの一覧を作る. + + Args: + root: プロジェクトルート． + + Returns: + arcname 順に並べた送信対象． + """ + return sorted(_walk_files(root.resolve()), key=lambda entry: entry.arcname) diff --git a/src/tias_judge_model/remote/gpu.py b/src/tias_judge_model/remote/gpu.py new file mode 100644 index 0000000..a8761ae --- /dev/null +++ b/src/tias_judge_model/remote/gpu.py @@ -0,0 +1,637 @@ +"""サーバの状況（tmux セッション・GPU・計算プロセス）の解析と，GPU の選択（純関数）. + +ssh を呼ぶ部分は CLI 側にあり，ここは ssh の出力文字列を受け取って解釈するだけにする +（テストで ssh・nvidia-smi をモックせず，出力の文字列で固定できるようにするため）． +判定と選択の規則は SPEC_03「GPU の空き判定と選択」に従う． +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from enum import StrEnum + +from tias_judge_model.remote.commands import ( + PROBE_APPS_MARKER, + PROBE_APPS_RC_PREFIX, + PROBE_END_MARKER, + PROBE_GPU_RC_PREFIX, + PROBE_GPUS_MARKER, + PROBE_JOB_PREFIX, + PROBE_OTHER_PREFIX, + PROBE_OWNERS_MARKER, + PROBE_PMON_MARKER, + PROBE_SAMPLES_MARKER, + PROBE_SESSIONS_MARKER, +) +from tias_judge_model.remote.config import GpuFreeThreshold, ServerConfig +from tias_judge_model.remote.observe import ( + GpuSample, + GpuUsage, + UsagePattern, + classify_usage_pattern, + parse_gpu_samples, + parse_pmon, + summarize_samples, +) + +# nvidia-smi --query-gpu の列数（index,uuid,name,memory.used,memory.total,utilization.gpu） +_GPU_COLUMNS = 6 +# nvidia-smi --query-compute-apps の列数（gpu_uuid,pid,used_memory） +_APP_COLUMNS = 3 + + +class GpuState(StrEnum): + """GPU の状態（選ぶ優先順は UNUSED → SHARED．BUSY・UNKNOWN は選ばない）.""" + + UNUSED = "未使用" + SHARED = "空き（利用者あり）" + BUSY = "使用中" + UNKNOWN = "状態不明" + + +# 自動選択で順に探す状態（未使用がサーバの優先順位より優先される） +SELECTABLE_STATES = (GpuState.UNUSED, GpuState.SHARED) + + +@dataclass(frozen=True) +class GpuProcess: + """GPU 上の計算プロセス 1 件. + + Attributes: + pid: プロセス ID． + used_memory_mib: 使用メモリ（MiB）．取れなければ None． + user: 所有ユーザー．ps で取れなければ None． + elapsed_sec: 起動からの経過秒数．ps で取れなければ None． + max_sm_percent: 観測中の SM 使用率の最大（pmon）．取れなければ None． + """ + + pid: int + used_memory_mib: int | None + user: str | None + elapsed_sec: int | None = None + max_sm_percent: int | None = None + + +@dataclass(frozen=True) +class ProcessOwner: + """ps で取ったプロセスの所有者と経過時間. + + Attributes: + user: 所有ユーザー． + elapsed_sec: 起動からの経過秒数．取れなければ None． + """ + + user: str + elapsed_sec: int | None + + +@dataclass(frozen=True) +class GpuInfo: + """GPU 1 基の状況. + + Attributes: + index: GPU 番号．解釈できなければ None． + uuid: GPU の UUID（計算プロセスとの対応付けに使う）．解釈できない行では空文字． + name: GPU の名前（解釈できない行では空文字）． + memory_used_mib: 使用中メモリ（MiB．観測前の瞬間値）．解釈できなければ None． + memory_total_mib: 総メモリ（MiB）．解釈できなければ None． + utilization_percent: 使用率（%．観測前の瞬間値）．解釈できなければ None． + state: 判定の結果． + processes: この GPU 上の計算プロセス． + processes_known: 計算プロセスの一覧が取れたか（False なら利用者不明）． + usage: 観測の要約（観測しないときは瞬間値 1 つの要約）．サンプルが無ければ None． + pattern: 使い方の型（表示専用）．決められなければ None． + """ + + index: int | None + uuid: str + name: str + memory_used_mib: int | None + memory_total_mib: int | None + utilization_percent: int | None + state: GpuState + processes: tuple[GpuProcess, ...] = () + processes_known: bool = False + usage: GpuUsage | None = None + pattern: UsagePattern | None = None + + +@dataclass(frozen=True) +class ProbeOutput: + """probe コマンド（`build_probe_command`）の出力を解釈した結果. + + Attributes: + jobs: このツールのジョブ（ジョブログがある tmux セッション）の名前． + other_sessions: それ以外の tmux セッションの名前． + gpu_query_ok: GPU の状態の取得が成功したか．False なら GPU はすべて候補から外す． + processes_query_ok: 計算プロセスの一覧の取得・解釈が成功したか．False なら閾値内の + GPU はすべて「空き（利用者あり）」になる． + gpus: GPU の一覧（GPU の状態の取得が失敗したときは空）． + observe_sec: 観測した秒数（0 なら瞬間値のみで判定した）． + pmon_ok: プロセス別の観測（pmon）が得られたか（False なら表示を省くだけ）． + """ + + jobs: tuple[str, ...] + other_sessions: tuple[str, ...] + gpu_query_ok: bool + processes_query_ok: bool + gpus: tuple[GpuInfo, ...] + observe_sec: int = 0 + pmon_ok: bool = False + + @property + def sessions(self) -> tuple[str, ...]: + """すべての tmux セッション名（ジョブとそれ以外）.""" + return self.jobs + self.other_sessions + + +@dataclass(frozen=True) +class ServerProbe: + """サーバ 1 台の状況. + + Attributes: + server: サーバの設定． + output: probe の結果．接続できなかったときは None． + error: 接続できなかった理由（表示用．ssh の標準エラーの最終行等）． + """ + + server: ServerConfig + output: ProbeOutput | None + error: str = "" + + @property + def reachable(self) -> bool: + """接続できたか.""" + return self.output is not None + + +@dataclass(frozen=True) +class GpuSelection: + """run が使う GPU. + + Attributes: + server: 選んだサーバ． + gpu_index: 選んだ GPU 番号． + state: 選んだ時点の GPU の状態． + gpu: 選んだ GPU の状況（利用者の警告表示用．状態を取れなかったときは None）． + forced: 使用中・状態不明の GPU を `--force` で選んだか（警告を出す）． + """ + + server: ServerConfig + gpu_index: int + state: GpuState + gpu: GpuInfo | None = None + forced: bool = False + + +def classify_gpu( + memory_used_mib: int, + utilization_percent: int, + threshold: GpuFreeThreshold, + process_count: int | None, +) -> GpuState: + """閾値と計算プロセスの数から GPU の状態を判定する. + + 観測するときは，使用中メモリ・使用率に**観測中の最大値**を渡す（一度でも閾値を超えた + GPU を使用中にするため）． + + Args: + memory_used_mib: 使用中メモリ（MiB．観測中の最大）． + utilization_percent: 使用率（%．観測中の最大）． + threshold: 空き判定の閾値． + process_count: 計算プロセスの数（所有者を問わない）．一覧が取れなければ None． + + Returns: + 閾値を満たさなければ `BUSY`．閾値内で計算プロセスが 0 なら `UNUSED`， + 1 つ以上または不明（None）なら `SHARED`（未使用と誤認しないため）． + """ + within = ( + memory_used_mib < threshold.max_memory_used_mib + and utilization_percent < threshold.max_utilization_percent + ) + if not within: + return GpuState.BUSY + return GpuState.UNUSED if process_count == 0 else GpuState.SHARED + + +def _parse_int(text: str) -> int | None: + """整数として解釈できれば返す（`[N/A]` 等は None）.""" + try: + return int(text.strip()) + except ValueError: + return None + + +def _is_valid_uuid(text: str) -> bool: + """nvidia-smi の UUID として使える値か（空・`[N/A]` 等を除く）.""" + return bool(text) and not text.startswith("[") and " " not in text + + +def parse_process_owners(text: str) -> dict[int, ProcessOwner]: + """probe の所有者の節（` <ユーザー> <経過秒>` の行）を解釈する. + + ユーザー名が無い行（ps で所有者を取れなかったプロセス）と解釈できない行は含めない + （所有者・経過時間は表示用で，判定には使わないため）． + + Args: + text: 所有者の節の出力． + + Returns: + pid から所有者・経過秒数への対応． + """ + owners: dict[int, ProcessOwner] = {} + for line in text.splitlines(): + tokens = line.split() + if len(tokens) < 2: + continue + pid = _parse_int(tokens[0]) + if pid is None: + continue + elapsed = _parse_int(tokens[2]) if len(tokens) >= 3 else None + owners[pid] = ProcessOwner(tokens[1], elapsed) + return owners + + +def parse_compute_apps( + text: str, + owners: Mapping[int, ProcessOwner], + sm_peaks: Mapping[int, int] | None = None, +) -> dict[str, tuple[GpuProcess, ...]] | None: + """`nvidia-smi --query-compute-apps=gpu_uuid,pid,used_memory` の出力を GPU ごとにまとめる. + + Args: + text: nvidia-smi の標準出力（空なら計算プロセスは 0）． + owners: pid から所有者・経過秒数への対応（`parse_process_owners`）． + sm_peaks: pid から観測中の SM 使用率の最大への対応（`parse_pmon`）．無ければ None． + + Returns: + GPU の UUID から計算プロセスへの対応．1 行でも解釈できなければ None + （一覧を信用せず利用者不明として扱うため）． + """ + apps: dict[str, list[GpuProcess]] = {} + for line in text.splitlines(): + if not line.strip(): + continue + columns = [column.strip() for column in line.split(",")] + if len(columns) != _APP_COLUMNS or not _is_valid_uuid(columns[0]): + return None + pid = _parse_int(columns[1]) + if pid is None: + return None + owner = owners.get(pid) + process = GpuProcess( + pid=pid, + used_memory_mib=_parse_int(columns[2]), + user=owner.user if owner else None, + elapsed_sec=owner.elapsed_sec if owner else None, + max_sm_percent=sm_peaks.get(pid) if sm_peaks is not None else None, + ) + apps.setdefault(columns[0], []).append(process) + return {uuid: tuple(processes) for uuid, processes in apps.items()} + + +def parse_nvidia_smi_line( + line: str, + threshold: GpuFreeThreshold, + processes: Mapping[str, Sequence[GpuProcess]] | None, + samples: Sequence[GpuSample] | None = None, +) -> GpuInfo: + """`nvidia-smi --query-gpu`（`--format=csv,noheader,nounits`）の 1 行を解釈して判定する. + + 列数不足・数値でない値（`[N/A]` 等）の行は `UNKNOWN` とする．GPU 名にカンマを含む場合に + 備え，先頭 2 列を index・uuid，末尾 3 列を数値として，その間を名前とする． + + 観測したとき（`samples` が None でない）は，観測中の最大値（とこの行の瞬間値の大きい方）で + 閾値を判定し，サンプルが 1 つも無ければ `UNKNOWN` とする．観測しないとき（None）は + この行の瞬間値で判定する． + + Args: + line: nvidia-smi の出力 1 行． + threshold: 空き判定の閾値． + processes: GPU の UUID から計算プロセスへの対応（`parse_compute_apps`）． + 計算プロセスの一覧が取れなかったときは None（閾値内なら `SHARED` になる）． + samples: この GPU の観測サンプル．観測しないときは None． + + Returns: + GPU の状況． + """ + columns = [column.strip() for column in line.split(",")] + if len(columns) < _GPU_COLUMNS: + index = _parse_int(columns[0]) if columns else None + return GpuInfo(index, "", "", None, None, None, GpuState.UNKNOWN) + index = _parse_int(columns[0]) + uuid = columns[1] + name = ", ".join(columns[2:-3]) + used, total, util = (_parse_int(value) for value in columns[-3:]) + # UUID が取れない GPU は計算プロセスと対応付けられないので，利用者不明として扱う + known = processes is not None and _is_valid_uuid(uuid) + gpu_processes = tuple(processes.get(uuid, ())) if known and processes is not None else () + count = len(gpu_processes) if known else None + + if index is None or used is None or total is None or util is None: + return GpuInfo(index, uuid, name, used, total, util, GpuState.UNKNOWN, gpu_processes, known) + observed = samples is not None + usage = summarize_samples(samples if samples is not None else [GpuSample(util, used)]) + if usage is None: + # 観測したのにサンプルが 1 つも得られない GPU は使わない + state = GpuState.UNKNOWN + pattern = None + else: + max_used = max(usage.max_memory_used_mib, used) if observed else used + max_util = max(usage.max_utilization_percent, util) if observed else util + state = classify_gpu(max_used, max_util, threshold, count) + pattern = classify_usage_pattern(usage, count) + return GpuInfo( + index, uuid, name, used, total, util, state, gpu_processes, known, usage, pattern + ) + + +def parse_nvidia_smi_output( + text: str, + threshold: GpuFreeThreshold, + processes: Mapping[str, Sequence[GpuProcess]] | None, + samples: Mapping[int, Sequence[GpuSample]] | None = None, +) -> list[GpuInfo]: + """`nvidia-smi --query-gpu` の出力全体を 1 行 1 GPU として解釈する（空行は無視する）. + + Args: + text: nvidia-smi の標準出力． + threshold: 空き判定の閾値． + processes: GPU の UUID から計算プロセスへの対応．一覧が取れなかったときは None． + samples: GPU 番号から観測サンプルへの対応（`parse_gpu_samples`）．観測しないときは + None（瞬間値で判定する）．観測したのに該当する GPU のサンプルが無ければ状態不明． + + Returns: + 出力順の GPU の一覧． + """ + gpus: list[GpuInfo] = [] + for line in text.splitlines(): + if not line.strip(): + continue + gpu_samples: Sequence[GpuSample] | None = None + if samples is not None: + index = _parse_int(line.split(",", 1)[0]) + gpu_samples = samples.get(index, []) if index is not None else [] + gpus.append(parse_nvidia_smi_line(line, threshold, processes, gpu_samples)) + return gpus + + +def parse_probe_output(text: str, threshold: GpuFreeThreshold) -> ProbeOutput: + """probe コマンド（`build_probe_command`）の標準出力を解釈する. + + 目印より前の行（ログインシェルが出す挨拶等）は無視する．計算プロセスの節が欠けている・ + 終了コードが 0 でない・解釈できない場合は，計算プロセスの一覧を「取得できない」とする． + `threshold.observe_sec > 0` なら観測したものとして判定し，観測の節が欠けていれば + （サーバで一時ディレクトリを作れなかった等）各 GPU は状態不明になる． + + Args: + text: probe コマンドの標準出力． + threshold: 空き判定の閾値と観測の設定（probe コマンドの組み立てに使ったもの）． + + Returns: + 解釈した結果． + + Raises: + ValueError: 必須の目印（セッション・GPU・GPU の終了コード・終端）がそろっていない + （出力が途中で切れた等）場合． + """ + section: str | None = None + jobs: list[str] = [] + others: list[str] = [] + lines: dict[str, list[str]] = { + "gpus": [], + "samples": [], + "pmon": [], + "apps": [], + "owners": [], + } + gpu_rc: int | None = None + apps_rc: int | None = None + seen_sessions = ended = False + for line in text.splitlines(): + if line == PROBE_SESSIONS_MARKER: + section, seen_sessions = "sessions", True + elif line == PROBE_GPUS_MARKER: + section = "gpus" + elif line.startswith(PROBE_GPU_RC_PREFIX) and section == "gpus": + gpu_rc = _parse_int(line[len(PROBE_GPU_RC_PREFIX) :]) + section = None + elif line == PROBE_SAMPLES_MARKER: + section = "samples" + elif line == PROBE_PMON_MARKER: + section = "pmon" + elif line == PROBE_APPS_MARKER: + section = "apps" + elif line.startswith(PROBE_APPS_RC_PREFIX) and section == "apps": + apps_rc = _parse_int(line[len(PROBE_APPS_RC_PREFIX) :]) + section = None + elif line == PROBE_OWNERS_MARKER: + section = "owners" + elif line == PROBE_END_MARKER: + section, ended = None, True + elif section == "sessions": + if line.startswith(PROBE_JOB_PREFIX): + jobs.append(line[len(PROBE_JOB_PREFIX) :]) + elif line.startswith(PROBE_OTHER_PREFIX): + others.append(line[len(PROBE_OTHER_PREFIX) :]) + elif section in lines: + lines[section].append(line) + if not (seen_sessions and ended) or gpu_rc is None: + raise ValueError("サーバの状況の出力を解釈できない（目印が欠けている）") + + texts = {name: "\n".join(section_lines) for name, section_lines in lines.items()} + gpu_ok = gpu_rc == 0 + observed = threshold.observe_sec > 0 + sm_peaks = parse_pmon(texts["pmon"]) if observed else None + processes = None + if gpu_ok and apps_rc == 0: + owners = parse_process_owners(texts["owners"]) + processes = parse_compute_apps(texts["apps"], owners, sm_peaks) + samples = parse_gpu_samples(texts["samples"]) if observed else None + gpus = ( + tuple(parse_nvidia_smi_output(texts["gpus"], threshold, processes, samples)) + if gpu_ok + else () + ) + return ProbeOutput( + jobs=tuple(jobs), + other_sessions=tuple(others), + gpu_query_ok=gpu_ok, + processes_query_ok=processes is not None, + gpus=gpus, + observe_sec=threshold.observe_sec if observed else 0, + pmon_ok=sm_peaks is not None, + ) + + +def _find_probe(probes: Sequence[ServerProbe], name: str) -> ServerProbe: + """名前（大文字小文字を区別しない）でサーバの状況を探す.""" + for probe in probes: + if probe.server.name.lower() == name.lower(): + return probe + known = ", ".join(probe.server.name for probe in probes) + raise ValueError(f"サーバ {name!r} は設定に無い（設定済み: {known}）") + + +def _select_explicit(probe: ServerProbe, gpu_index: int, *, force: bool) -> GpuSelection: + """`--server --gpu` で指定された GPU を検証して返す.""" + output = probe.output + if output is None: + raise RuntimeError("接続できないサーバの GPU を検証しようとした") + gpu = next((g for g in output.gpus if g.index == gpu_index), None) + if gpu is None: + # 一覧を正しく取れていて該当が無いなら，存在しない GPU の指定である + if output.gpu_query_ok and all(g.index is not None for g in output.gpus): + raise ValueError(f"{probe.server.name} に GPU {gpu_index} は無い") + state = GpuState.UNKNOWN + else: + state = gpu.state + if state in SELECTABLE_STATES: + return GpuSelection(probe.server, gpu_index, state, gpu) + if not force: + raise ValueError( + f"{probe.server.name} の GPU {gpu_index} は{state}である" + "（使うなら --force を付ける．他のユーザーが使用中の可能性がある）" + ) + return GpuSelection(probe.server, gpu_index, state, gpu, forced=True) + + +def check_selection_options(*, server_name: str | None, gpu_index: int | None, force: bool) -> None: + """GPU 指定の組み合わせを検証する（サーバへ問い合わせる前に弾けるよう独立させている）. + + Args: + server_name: `--server` の値． + gpu_index: `--gpu` の値． + force: `--force` の有無． + + Raises: + ValueError: `--gpu` だけの指定，`--gpu` 無しの `--force`，負の GPU 番号の場合． + """ + if gpu_index is not None and server_name is None: + raise ValueError("--gpu を指定するときは --server も指定する（どのサーバか曖昧になるため）") + if force and gpu_index is None: + raise ValueError("--force は --server と --gpu で GPU を指定したときだけ使える") + if gpu_index is not None and gpu_index < 0: + raise ValueError(f"GPU 番号は 0 以上: {gpu_index}") + + +def select_gpu( + probes: Sequence[ServerProbe], + *, + server_name: str | None = None, + gpu_index: int | None = None, + force: bool = False, + include_not_ready: bool = False, +) -> GpuSelection | None: + """run が使う GPU を選ぶ（SPEC_03「GPU の空き判定と選択」）. + + 自動選択では，候補サーバ（`data_ready` のもの．`include_not_ready` なら全サーバ． + `server_name` を指定したらそのサーバのみ）全体を通して，まず「未使用」を priority 昇順 + （同値は `probes` の順）→ GPU 番号の昇順で探し，無ければ同じ順で「空き（利用者あり）」を + 探す（未使用がサーバの優先順位より優先される）．接続できないサーバと GPU の状態の取得に + 失敗したサーバは候補から外す． + + Args: + probes: 各サーバの状況（設定の記載順で渡す．同じ priority の順位に使う）． + server_name: `--server` の値． + gpu_index: `--gpu` の値（`server_name` が必須）． + force: 指定した GPU が使用中・状態不明でも選ぶ（`gpu_index` 指定時のみ有効）． + include_not_ready: `data_ready: false` のサーバも候補にする（`--no-data`）． + + Returns: + 選んだ GPU．自動選択で未使用・空き（利用者あり）のどちらも無ければ None． + + Raises: + ValueError: 指定の組み合わせが不正（`--gpu` だけ・`--force` だけ），指定したサーバが + 設定に無い・データ未配置・接続不可，指定した GPU が無い・使用中・状態不明 + （`force` 無し）の場合． + """ + check_selection_options(server_name=server_name, gpu_index=gpu_index, force=force) + if server_name is not None: + probe = _find_probe(probes, server_name) + if not probe.server.data_ready and not include_not_ready: + raise ValueError( + f"{probe.server.name} は臨床データが未配置（data_ready: false）である" + "（データを使わないジョブなら --no-data を付ける）" + ) + if not probe.reachable: + raise ValueError(f"{probe.server.name} に接続できない（{probe.error}）") + if gpu_index is not None: + return _select_explicit(probe, gpu_index, force=force) + candidates = [probe] + else: + candidates = sorted( + (p for p in probes if p.reachable and (include_not_ready or p.server.data_ready)), + key=lambda p: p.server.priority, + ) + + for wanted in SELECTABLE_STATES: + for probe in candidates: + output = probe.output + if output is None or not output.gpu_query_ok: + continue + for gpu in sorted( + (g for g in output.gpus if g.index is not None), key=lambda g: g.index or 0 + ): + if gpu.state is wanted and gpu.index is not None: + return GpuSelection(probe.server, gpu.index, gpu.state, gpu) + return None + + +def display_user_name(user: str | None, own_user: str) -> str: + """計算プロセスの所有ユーザーを表示用の名前にする. + + Args: + user: 所有ユーザー．ps で取れなければ None． + own_user: 自分のユーザー名（ssh の接続ユーザー）． + + Returns: + 自分なら「自分」，取れなければ「不明」，それ以外はユーザー名． + """ + if user is None: + return "不明" + if user == own_user: + return "自分" + return user + + +def gpu_users(gpu: GpuInfo, own_user: str) -> list[str]: + """GPU の計算プロセスの所有ユーザーを表示用に並べる（重複は 1 つにまとめ，出現順）. + + Args: + gpu: GPU の状況． + own_user: 自分のユーザー名（ssh の接続ユーザー）．一致するものは「自分」と表示する． + + Returns: + ユーザー名の一覧（自分は「自分」，所有者を取れなかったものは「不明」）． + """ + names: list[str] = [] + for process in gpu.processes: + name = display_user_name(process.user, own_user) + if name not in names: + names.append(name) + return names + + +def servers_with_session( + probes: Sequence[ServerProbe], name: str, *, jobs_only: bool +) -> list[ServerProbe]: + """指定した名前の tmux セッションがあるサーバを返す. + + Args: + probes: 各サーバの状況（接続できなかったサーバは対象外）． + name: セッション（ジョブ）名． + jobs_only: True ならこのツールのジョブ（ジョブログがあるもの）だけを見る． + + Returns: + 該当するサーバの状況（`probes` の順）． + """ + found: list[ServerProbe] = [] + for probe in probes: + if probe.output is None: + continue + names = probe.output.jobs if jobs_only else probe.output.sessions + if name in names: + found.append(probe) + return found diff --git a/src/tias_judge_model/remote/observe.py b/src/tias_judge_model/remote/observe.py new file mode 100644 index 0000000..6be3fac --- /dev/null +++ b/src/tias_judge_model/remote/observe.py @@ -0,0 +1,188 @@ +"""GPU の観測結果（サンプリング・pmon）の解析と，使い方の型の判定（純関数）. + +観測は probe コマンドがサーバ側で行い（`build_probe_command`），ここはその出力文字列を +解釈するだけにする．使い方の型は表示専用で，空き判定には使わない（SPEC_03「使い方の型」）． +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from enum import StrEnum + +# 0% を超えたサンプルがこの割合以上なら「常時使用」 +CONSTANT_USE_RATIO = 0.8 +# nvidia-smi --query-gpu=index,utilization.gpu,memory.used の列数 +_SAMPLE_COLUMNS = 3 + + +class UsagePattern(StrEnum): + """観測中の使用率から付ける使い方の型（表示専用）.""" + + CONSTANT = "常時使用" + INTERMITTENT = "断続的に使用" + MEMORY_ONLY = "メモリ確保のみ" + IDLE = "未使用" + + +@dataclass(frozen=True) +class GpuSample: + """GPU 1 基の 1 回分のサンプル. + + Attributes: + utilization_percent: 使用率（%）． + memory_used_mib: 使用中メモリ（MiB）． + """ + + utilization_percent: int + memory_used_mib: int + + +@dataclass(frozen=True) +class GpuUsage: + """観測中のサンプルの要約. + + Attributes: + sample_count: サンプル数（1 以上）． + active_count: 使用率が 0% を超えたサンプル数． + mean_utilization_percent: 使用率の平均（%）． + max_utilization_percent: 使用率の最大（%）． + max_memory_used_mib: 使用中メモリの最大（MiB）． + """ + + sample_count: int + active_count: int + mean_utilization_percent: float + max_utilization_percent: int + max_memory_used_mib: int + + +def _parse_int(text: str) -> int | None: + """整数として解釈できれば返す（`[N/A]`・`-` 等は None）.""" + try: + return int(text.strip()) + except ValueError: + return None + + +def parse_gpu_samples(text: str) -> dict[int, list[GpuSample]]: + """観測のサンプリング出力（`index, 使用率, 使用中メモリ` の CSV）を GPU ごとにまとめる. + + `timeout` で打ち切った最後の行が途中で切れていることがあるため，解釈できない行は + 読み飛ばす（そのぶんサンプルが減るだけで，サンプルが 0 件の GPU は状態不明になる）． + + Args: + text: サンプリングの出力． + + Returns: + GPU 番号からサンプル（出力順）への対応． + """ + samples: dict[int, list[GpuSample]] = {} + for line in text.splitlines(): + columns = line.split(",") + if len(columns) != _SAMPLE_COLUMNS: + continue + index, util, memory = (_parse_int(column) for column in columns) + if index is None or util is None or memory is None: + continue + samples.setdefault(index, []).append(GpuSample(util, memory)) + return samples + + +def summarize_samples(samples: Sequence[GpuSample]) -> GpuUsage | None: + """サンプルを要約する. + + Args: + samples: GPU 1 基のサンプル． + + Returns: + 要約．サンプルが無ければ None． + """ + if not samples: + return None + utils = [sample.utilization_percent for sample in samples] + return GpuUsage( + sample_count=len(samples), + active_count=sum(1 for util in utils if util > 0), + mean_utilization_percent=sum(utils) / len(utils), + max_utilization_percent=max(utils), + max_memory_used_mib=max(sample.memory_used_mib for sample in samples), + ) + + +def classify_usage_pattern(usage: GpuUsage, process_count: int | None) -> UsagePattern | None: + """観測中の使用率から使い方の型を付ける（表示専用）. + + Args: + usage: 観測の要約． + process_count: 計算プロセスの数．一覧が取れなければ None． + + Returns: + 使い方の型．0% を超えたサンプルが無く，計算プロセスの有無も分からなければ None． + """ + if usage.active_count > 0: + if usage.active_count >= CONSTANT_USE_RATIO * usage.sample_count: + return UsagePattern.CONSTANT + return UsagePattern.INTERMITTENT + if process_count is None: + return None + return UsagePattern.MEMORY_ONLY if process_count > 0 else UsagePattern.IDLE + + +def parse_pmon(text: str) -> dict[int, int] | None: + """`nvidia-smi pmon -s u` の出力から，計算プロセスごとの SM 使用率の最大を求める. + + `#` で始まる行はヘッダ（1 行目が列名 `# gpu pid type sm ...`，2 行目が単位）．列の位置は + ヘッダの列名から決める．種別に `C`（計算）を含む行だけを数え，`G`（描画）だけの行は除く． + 値が `-` の行（その回はプロセスが無い・値が取れない）は読み飛ばす． + + Args: + text: pmon の出力． + + Returns: + pid から SM 使用率の最大（%）への対応．列名のヘッダが無い（pmon が失敗・未対応） + ときは None（表示を省くだけで，判定には影響させない）． + """ + columns: list[str] | None = None + peaks: dict[int, int] = {} + for line in text.splitlines(): + stripped = line.strip() + if not stripped: + continue + if stripped.startswith("#"): + names = stripped.lstrip("#").split() + if columns is None and {"pid", "type", "sm"} <= set(names): + columns = names + continue + if columns is None: + continue + values = stripped.split() + if len(values) < len(columns) - 1: + continue + row = dict(zip(columns, values, strict=False)) + pid = _parse_int(row.get("pid", "-")) + sm = _parse_int(row.get("sm", "-")) + if pid is None or sm is None or "C" not in row.get("type", ""): + continue + peaks[pid] = max(peaks.get(pid, 0), sm) + return None if columns is None else peaks + + +def format_elapsed(seconds: int) -> str: + """経過秒数を読みやすい単位で返す（例: `45 秒`・`12.5 分`・`27.7 時間`・`3.2 日`）. + + 48 時間までは時間で示す（日単位にすると 1 日前後の違いが読み取りにくいため）． + + Args: + seconds: 経過秒数（0 以上）． + + Returns: + 表示用の文字列． + """ + if seconds < 60: + return f"{seconds} 秒" + if seconds < 3600: + return f"{seconds / 60:.1f} 分" + if seconds < 48 * 3600: + return f"{seconds / 3600:.1f} 時間" + return f"{seconds / 86400:.1f} 日" diff --git a/src/tias_judge_model/remote/report.py b/src/tias_judge_model/remote/report.py new file mode 100644 index 0000000..3a6f636 --- /dev/null +++ b/src/tias_judge_model/remote/report.py @@ -0,0 +1,193 @@ +"""サーバの状況（GPU 一覧・実行中ジョブ）を表示用の文字列に整形する純関数群.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from tias_judge_model.remote.gpu import ( + GpuInfo, + GpuProcess, + GpuSelection, + GpuState, + ServerProbe, + display_user_name, + gpu_users, +) +from tias_judge_model.remote.observe import format_elapsed + +# 計算プロセスの一覧が取れなかった GPU の利用者の表示（run の警告と GPU 一覧で共通） +_UNKNOWN_USERS = "利用者不明（計算プロセスの一覧を取得できない）" + + +def _format_value(value: int | None, unit: str) -> str: + return "?" if value is None else f"{value}{unit}" + + +def describe_gpu_users(gpu: GpuInfo, own_user: str) -> str: + """GPU の利用者（計算プロセスの所有ユーザー）を表示用に 1 句で返す. + + Args: + gpu: GPU の状況． + own_user: 自分のユーザー名（ssh の接続ユーザー．一致するものは「自分」）． + + Returns: + `利用者: 自分, alice` の形．計算プロセスの一覧が取れなければ `利用者不明`， + 計算プロセスが無ければ空文字． + """ + if not gpu.processes_known: + return _UNKNOWN_USERS + users = gpu_users(gpu, own_user) + return f"利用者: {', '.join(users)}" if users else "" + + +def describe_usage(gpu: GpuInfo, *, observed: bool) -> str: + """GPU の使い方（型と使用率）を表示用に 1 句で返す. + + Args: + gpu: GPU の状況． + observed: 観測したか（False なら瞬間値だけを示し，型は付けない）． + + Returns: + 観測したときは `断続的に使用・平均 21%・最大 74%（0% 超 7/20）` の形．観測サンプルが + 無ければ `観測できない`．観測しないときは `使用率 29%（瞬間値）`． + """ + usage = gpu.usage + if not observed: + return f"使用率 {_format_value(gpu.utilization_percent, '%')}（瞬間値）" + if usage is None: + return "観測できない（サンプル無し）" + pattern = gpu.pattern.value if gpu.pattern is not None else "使い方不明" + return ( + f"{pattern}・平均 {usage.mean_utilization_percent:.0f}%" + f"・最大 {usage.max_utilization_percent}%" + f"（0% 超 {usage.active_count}/{usage.sample_count}）" + ) + + +def describe_process(process: GpuProcess, own_user: str) -> str: + """計算プロセス 1 件を表示用に 1 句で返す. + + Args: + process: 計算プロセス． + own_user: 自分のユーザー名（ssh の接続ユーザー．一致するものは「自分」）． + + Returns: + `ando 経過 27.7 時間 SM 最大 74%` の形（取れない項目は省く．所有者が取れなければ + `不明`）． + """ + parts = [display_user_name(process.user, own_user)] + if process.elapsed_sec is not None: + parts.append(f"経過 {format_elapsed(process.elapsed_sec)}") + if process.max_sm_percent is not None: + parts.append(f"SM 最大 {process.max_sm_percent}%") + return " ".join(parts) + + +def _format_gpu_lines(gpu: GpuInfo, own_user: str, *, observed: bool) -> list[str]: + index = "?" if gpu.index is None else str(gpu.index) + memory = f"{_format_value(gpu.memory_used_mib, '')} / {_format_value(gpu.memory_total_mib, '')}" + lines = [ + f" GPU {index:>2} {gpu.name} メモリ {memory} MiB [{gpu.state.value}]" + f" {describe_usage(gpu, observed=observed)}" + ] + if gpu.state is GpuState.UNKNOWN and gpu.usage is None and not gpu.processes: + return lines + if not gpu.processes_known: + lines.append(f" {_UNKNOWN_USERS}") + else: + lines.extend(f" - {describe_process(p, own_user)}" for p in gpu.processes) + return lines + + +def format_server_header(probe: ServerProbe) -> str: + """サーバ 1 台の見出し行を返す. + + Args: + probe: サーバの状況． + + Returns: + `[名前] アドレス priority=N データ=... 観測 10 秒` の形の 1 行（接続不可なら理由を + 添える）． + """ + server = probe.server + data = "配置済み" if server.data_ready else "未配置" + header = f"[{server.name}] {server.address} priority={server.priority} データ={data}" + if probe.output is None: + reason = f"（{probe.error}）" if probe.error else "" + return f"{header} 接続不可{reason}" + observe_sec = probe.output.observe_sec + return f"{header} {'観測 ' + str(observe_sec) + ' 秒' if observe_sec else '瞬間値のみ'}" + + +def format_server_report(probes: Sequence[ServerProbe], *, show_jobs: bool) -> str: + """全サーバの GPU 一覧（状態と利用者．と実行中ジョブ）を表示用に整形する. + + Args: + probes: 各サーバの状況（表示順）． + show_jobs: 実行中ジョブも表示するか（status 用）． + + Returns: + 末尾に改行を含む複数行の文字列． + """ + lines: list[str] = [] + for probe in probes: + lines.append(format_server_header(probe)) + output = probe.output + if output is None: + continue + if show_jobs: + jobs = ", ".join(output.jobs) if output.jobs else "（なし）" + lines.append(f" 実行中のジョブ: {jobs}") + if output.other_sessions: + lines.append( + f" （このツール以外の tmux セッション {len(output.other_sessions)} 件）" + ) + if not output.gpu_query_ok: + lines.append(" GPU: nvidia-smi が失敗した（このサーバの GPU は候補から外す）") + elif not output.gpus: + lines.append(" GPU: なし") + else: + lines.append(" GPU:") + for gpu in output.gpus: + lines.extend( + _format_gpu_lines(gpu, probe.server.user, observed=output.observe_sec > 0) + ) + return "\n".join(lines) + "\n" + + +def format_selection(selection: GpuSelection | None) -> str: + """run が選ぶ GPU の説明を 1 行で返す. + + Args: + selection: `select_gpu` の結果． + + Returns: + 表示用の 1 行（改行なし）． + """ + if selection is None: + return "run が選ぶ GPU: なし（未使用・空き（利用者あり）の GPU が無い）" + text = ( + f"run が選ぶ GPU: {selection.server.name} の GPU {selection.gpu_index}" + f"（{selection.state.value}）" + ) + if selection.state is GpuState.SHARED and selection.gpu is not None: + text += f" {describe_gpu_users(selection.gpu, selection.server.user)}" + return text + + +def format_shared_warning(selection: GpuSelection) -> str | None: + """「空き（利用者あり）」の GPU を選んだときの警告文を返す（それ以外は None）. + + Args: + selection: `select_gpu` の結果． + + Returns: + 警告文（改行なし）． + """ + if selection.state is not GpuState.SHARED or selection.gpu is None: + return None + users = describe_gpu_users(selection.gpu, selection.server.user) + return ( + f"{selection.server.name} の GPU {selection.gpu_index} には計算プロセスがある" + f"（{users}）．他の計算とメモリ・計算時間を分け合う" + ) diff --git a/src/tias_judge_model/remote/transfer.py b/src/tias_judge_model/remote/transfer.py new file mode 100644 index 0000000..a957329 --- /dev/null +++ b/src/tias_judge_model/remote/transfer.py @@ -0,0 +1,121 @@ +"""tar ストリームの書き出し（push）と，回収した tar の検査・展開（pull）. + +Windows には rsync が無いため，Python の tarfile でストリームを作り `ssh ... 'tar -x ...'` の +標準入力へ流す（push）．pull は ssh の標準出力をいったん一時ファイルに受け取り， +全メンバーを検査してから展開する（ストリームのまま展開すると，問題のメンバーに当たる前の +ファイルが展開済みになってしまうため．SPEC_03「pull」の二重の防御）． +""" + +from __future__ import annotations + +import tarfile +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import BinaryIO + +from tias_judge_model.remote.files import PushEntry +from tias_judge_model.remote.whitelist import check_pull_members + +_ENCODING = "utf-8" +_COPY_CHUNK = 1024 * 1024 + + +@dataclass(frozen=True) +class ExtractResult: + """展開の結果. + + Attributes: + files: 展開したファイル数． + bytes: 展開したファイルの合計バイト数． + """ + + files: int + bytes: int + + +def _normalize_tarinfo(info: tarfile.TarInfo) -> tarfile.TarInfo: + """所有者とパーミッションを送信元 OS に依存しない値にそろえる. + + Windows の stat は所有者を持たず実行ビットも拡張子から推測した値なので， + 展開結果がそれに引きずられないよう固定する（実行ビットが立っているものだけ 755）． + """ + info.uid = info.gid = 0 + info.uname = info.gname = "" + info.mode = 0o755 if info.isdir() or info.mode & 0o111 else 0o644 + return info + + +def write_tar_stream(entries: Sequence[PushEntry], fileobj: BinaryIO) -> int: + """送信対象を非圧縮の tar ストリームとして fileobj へ書き出す. + + PAX 形式で書く（日本語ファイル名を UTF-8 で確実に運ぶ）．fileobj は閉じない + （呼び出し側が ssh の標準入力を閉じて EOF を伝える）． + + Args: + entries: 送信対象． + fileobj: 書き込み先のバイナリストリーム（ssh の標準入力）． + + Returns: + 書き出したファイルの合計バイト数． + + Raises: + OSError: 書き込み先が閉じられた（ssh が異常終了した）場合など． + """ + done = 0 + with tarfile.open( + fileobj=fileobj, mode="w|", format=tarfile.PAX_FORMAT, encoding=_ENCODING, dereference=True + ) as tar: + for entry in entries: + tar.add(entry.path, arcname=entry.arcname, recursive=False, filter=_normalize_tarinfo) + done += entry.size + return done + + +def copy_stream(source: BinaryIO, destination: BinaryIO) -> int: + """source を終端まで読み，destination へ書き写す. + + Args: + source: 読み出し元（ssh の標準出力）． + destination: 書き込み先（一時ファイル）． + + Returns: + 書き写したバイト数． + """ + total = 0 + while chunk := source.read(_COPY_CHUNK): + destination.write(chunk) + total += len(chunk) + return total + + +def extract_pull_archive( + archive: BinaryIO, dest: Path, run_id: str, *, with_checkpoints: bool +) -> ExtractResult: + """回収した tar（シーク可能なファイル）の全メンバーを検査し，問題が無ければ dest へ展開する. + + 検査（`check_pull_members`）で 1 つでも問題があれば，何も展開せずに例外を送出する． + 検査を通ったメンバーは相対パスの通常ファイルのみだが，念のため Python 3.12 の + `filter="data"`（展開先の外・特殊ファイル・危険なパーミッションを拒否）も重ねる． + + Args: + archive: tar を書き込んだ一時ファイル（読み出し位置は問わない）． + dest: 展開先（手元の `runs/`．`/...` がその下にできる．同名ファイルは上書き）． + run_id: 検証済みの RUN_ID． + with_checkpoints: `checkpoints/*.pt` を受け入れるか． + + Returns: + 展開したファイル数と合計バイト数． + + Raises: + ValueError: 回収できないメンバーが含まれる場合（何も展開しない）． + tarfile.TarError: アーカイブが壊れている場合（何も展開しない）． + OSError: 書き込みに失敗した場合． + """ + archive.seek(0) + with tarfile.open(fileobj=archive, mode="r:", encoding=_ENCODING) as tar: + members = tar.getmembers() + check_pull_members(members, run_id, with_checkpoints=with_checkpoints) + dest.mkdir(parents=True, exist_ok=True) + tar.extractall(dest, members=members, filter="data") + return ExtractResult(files=len(members), bytes=sum(member.size for member in members)) diff --git a/src/tias_judge_model/remote/whitelist.py b/src/tias_judge_model/remote/whitelist.py new file mode 100644 index 0000000..aa2da38 --- /dev/null +++ b/src/tias_judge_model/remote/whitelist.py @@ -0,0 +1,215 @@ +"""pull の回収ホワイトリスト（SPEC_01「回収ホワイトリスト」）の判定と，受け取った tar の検査. + +回収できるのは `runs//` 配下で名前が一致するファイルのみ．拡張子や中身では判断せず， +列挙外はすべて回収しない．判定は 2 か所で使う（二重の防御．SPEC_03「pull」）． + +1. サーバの一覧（`build_pull_list_command` の出力）から送らせるファイルを絞り込む（`plan_pull`） +2. 手元で受け取った tar を展開前に全メンバー検査する（`check_pull_members`） + +エラーメッセージ・戻り値の表示用情報には件数だけを載せ，ファイル名は出さない +（患者情報を含む名前の可能性があるため）． +""" + +from __future__ import annotations + +import re +import tarfile +from collections.abc import Iterable +from dataclasses import dataclass + +# runs// 直下で回収するファイル名 +PULL_ROOT_FILES = frozenset( + {"config.yaml", "history.csv", "metrics.json", "confusion_matrix.csv", "log.txt"} +) +PULL_FIGURES_DIR = "figures" +PULL_FIGURE_SUFFIX = ".png" +PULL_CHECKPOINTS_DIR = "checkpoints" +PULL_CHECKPOINT_SUFFIX = ".pt" + +# find -printf '%y' の通常ファイル +_FIND_TYPE_FILE = "f" +_WINDOWS_DRIVE_RE = re.compile(r"^[A-Za-z]:") + + +@dataclass(frozen=True) +class PullPlan: + """サーバから送らせるファイルの計画. + + Attributes: + files: 送らせるファイル（`runs/` 基準の POSIX パス．表示しない）． + skipped: 回収しなかったファイルの件数（リンク等の通常ファイル以外を含む）． + skipped_checkpoints: そのうち `--with-checkpoints` を付ければ回収できた件数． + """ + + files: tuple[str, ...] + skipped: int + skipped_checkpoints: int + + +def validate_run_id(run_id: str) -> str: + """RUN_ID（`runs/` 直下のディレクトリ名）を検証して返す. + + Args: + run_id: 実験 ID． + + Returns: + 検証済みの RUN_ID． + + Raises: + ValueError: 空・`/`・`\\`・`..`・制御文字を含む，`.`/`-` で始まる場合． + """ + if ( + not run_id + or "/" in run_id + or "\\" in run_id + or ".." in run_id + or run_id.startswith((".", "-")) + or any(ord(c) < 0x20 or c == "\x7f" for c in run_id) + or run_id != run_id.strip() + ): + raise ValueError( + f"RUN_ID が不正である: {run_id!r}（runs/ 直下のディレクトリ名を指定する．" + "/・\\・.. は使えない）" + ) + return run_id + + +def _is_safe_part(part: str) -> bool: + """パスの 1 要素が `..`・`.`・空・バックスラッシュ・ドライブ文字を含まないか.""" + return part not in ("", ".", "..") and "\\" not in part and not _WINDOWS_DRIVE_RE.match(part) + + +def _is_safe_relative(path: str) -> bool: + """POSIX 区切りの相対パスとして安全か（絶対パス・パス外へ出る要素を含まない）.""" + return not path.startswith("/") and all(_is_safe_part(part) for part in path.split("/")) + + +def is_pull_allowed(rel_path: str, *, with_checkpoints: bool) -> bool: + """`runs//` 基準の相対パスが回収ホワイトリストに一致するかを返す. + + Args: + rel_path: `runs//` 基準の POSIX 相対パス． + with_checkpoints: `checkpoints/*.pt` も回収対象にするか． + + Returns: + 回収してよければ True． + """ + if not _is_safe_relative(rel_path): + return False + parts = rel_path.split("/") + if len(parts) == 1: + return parts[0] in PULL_ROOT_FILES + if len(parts) != 2: + return False + directory, name = parts + if directory == PULL_FIGURES_DIR: + return name.endswith(PULL_FIGURE_SUFFIX) and len(name) > len(PULL_FIGURE_SUFFIX) + if directory == PULL_CHECKPOINTS_DIR and with_checkpoints: + return name.endswith(PULL_CHECKPOINT_SUFFIX) and len(name) > len(PULL_CHECKPOINT_SUFFIX) + return False + + +def _strip_run_prefix(path: str, run_id: str) -> str | None: + """`/` から `` を取り出す（RUN_ID 配下でなければ None）.""" + prefix = f"{run_id}/" + return path[len(prefix) :] if path.startswith(prefix) else None + + +def parse_find_listing(data: bytes) -> list[tuple[str, str]]: + """`build_pull_list_command` の出力（`<種類> <パス>` の NUL 区切り）を分解する. + + Linux のファイル名は任意のバイト列なので，往復で失われないよう surrogateescape で扱う． + + Args: + data: サーバの出力． + + Returns: + (種類, `runs/` 基準のパス) の一覧． + + Raises: + ValueError: 書式が崩れている場合（件数のみ示す）． + """ + entries: list[tuple[str, str]] = [] + for number, record in enumerate((r for r in data.split(b"\0") if r), start=1): + kind, sep, path = record.decode("utf-8", "surrogateescape").partition(" ") + if not sep or not kind or not path: + raise ValueError(f"サーバのファイル一覧の {number} 件目を解釈できない") + entries.append((kind, path)) + return entries + + +def plan_pull( + listing: Iterable[tuple[str, str]], run_id: str, *, with_checkpoints: bool +) -> PullPlan: + """サーバのファイル一覧から，送らせるファイル（ホワイトリストに一致する通常ファイル）を選ぶ. + + Args: + listing: `parse_find_listing` の結果． + run_id: 検証済みの RUN_ID． + with_checkpoints: `checkpoints/*.pt` も回収するか． + + Returns: + 送らせるファイルと，回収しなかった件数． + """ + files: list[str] = [] + skipped = 0 + skipped_checkpoints = 0 + for kind, path in listing: + rel = _strip_run_prefix(path, run_id) + if kind == _FIND_TYPE_FILE and rel is not None: + if is_pull_allowed(rel, with_checkpoints=with_checkpoints): + files.append(path) + continue + if not with_checkpoints and is_pull_allowed(rel, with_checkpoints=True): + skipped_checkpoints += 1 + skipped += 1 + return PullPlan(tuple(sorted(files)), skipped, skipped_checkpoints) + + +def check_pull_members( + members: Iterable[tarfile.TarInfo], run_id: str, *, with_checkpoints: bool +) -> int: + """受け取った tar の全メンバーを展開前に検査する. + + 1 つでも問題があれば例外を送出する（呼び出し側は何も展開しない）．受け入れるのは + `/<ホワイトリストに一致する名前>` の通常ファイルのみで，ディレクトリ・リンク + （シンボリック・ハード）・特殊ファイルは拒否する（サーバはファイル名だけを tar に渡すため， + 正常な応答にディレクトリのメンバーは含まれない）． + + Args: + members: tar のメンバー（`TarFile.getmembers()`）． + run_id: 検証済みの RUN_ID． + with_checkpoints: `checkpoints/*.pt` を受け入れるか． + + Returns: + 受け入れたメンバーの件数． + + Raises: + ValueError: ホワイトリスト外の名前・パス外に出る名前（`..`・絶対パス・ドライブ文字・ + バックスラッシュ）・リンク・通常ファイル以外が含まれる場合（件数のみを示す）． + """ + accepted = 0 + outside = 0 + links = 0 + not_regular = 0 + not_allowed = 0 + for member in members: + if not _is_safe_relative(member.name): + outside += 1 + elif member.issym() or member.islnk(): + links += 1 + elif not member.isreg(): + not_regular += 1 + else: + rel = _strip_run_prefix(member.name, run_id) + if rel is None or not is_pull_allowed(rel, with_checkpoints=with_checkpoints): + not_allowed += 1 + else: + accepted += 1 + if outside or links or not_regular or not_allowed: + raise ValueError( + "受け取ったアーカイブに回収できないメンバーが含まれるため，何も展開せず中止した" + f"（ホワイトリスト外 {not_allowed} 件・パス外 {outside} 件・リンク {links} 件・" + f"通常ファイル以外 {not_regular} 件）" + ) + return accepted diff --git a/tests/cli/test_remote.py b/tests/cli/test_remote.py new file mode 100644 index 0000000..f911a62 --- /dev/null +++ b/tests/cli/test_remote.py @@ -0,0 +1,457 @@ +"""tias-remote CLI のテスト（SPEC_03）．ssh はすべてモックし，実サーバに接続しない. + +- 接続できないサーバがあっても他のサーバで続行する（全台不可ならエラー） +- `TIAS_REMOTE_JUMP` が空文字なら踏み台を使わない +- 手元のホスト名がサーバ名と一致したら止まる（`TIAS_REMOTE_ALLOW_SELF=1` で無効化） +- pull は回収ホワイトリストに一致するファイルだけを持ち帰り，回収しなかったファイルの名前を + 出さない．受け取った tar に不正なメンバーがあれば何も展開しない +""" + +from __future__ import annotations + +import io +import logging +import tarfile +import threading +from collections.abc import Callable +from pathlib import Path + +import pytest + +import tias_judge_model.cli.remote as cli +from tias_judge_model.remote.commands import ( + EXIT_NOT_FOUND, + PROBE_APPS_MARKER, + PROBE_APPS_RC_PREFIX, + PROBE_END_MARKER, + PROBE_GPU_RC_PREFIX, + PROBE_GPUS_MARKER, + PROBE_OWNERS_MARKER, + PROBE_SAMPLES_MARKER, + PROBE_SESSIONS_MARKER, +) +from tias_judge_model.remote.gpu import GpuState + +NEWTON_ADDR = "192.168.23.14" +ARCHIMEDES_ADDR = "192.168.23.15" +RUN_ID = "20260927-120000_newton_abc1234" + +CONFIG_YAML = f"""\ +jump_host: labpc +connect_timeout_sec: 15 +gpu_free: + max_memory_used_mib: 1024 + max_utilization_percent: 10 + observe_sec: 10 + sample_interval_ms: 500 +servers: + - name: Newton + address: {NEWTON_ADDR} + user: hasegawa + priority: 1 + remote_dir: ~/work/TIASJudgeModel + data_dir: ~/data/TIASJudgeModel + uv: ~/.local/bin/uv + data_ready: true + - name: Archimedes + address: {ARCHIMEDES_ADDR} + user: hasegawa + priority: 2 + remote_dir: ~/work/TIASJudgeModel + data_dir: ~/data/TIASJudgeModel + uv: ~/.local/bin/uv + data_ready: true +""" + + +def _probe_stdout(gpus: list[tuple[int, int, int]], apps: list[str] | None = None) -> bytes: + """probe の出力（gpus は (index, 使用中メモリ, 使用率) のリスト．観測サンプルも同じ値）.""" + lines = [PROBE_SESSIONS_MARKER, PROBE_GPUS_MARKER] + lines += [f"{i}, GPU-{i}, NVIDIA RTX A6000, {mem}, 49140, {util}" for i, mem, util in gpus] + lines.append(f"{PROBE_GPU_RC_PREFIX}0") + lines.append(PROBE_SAMPLES_MARKER) + lines += [f"{i}, {util}, {mem}" for i, mem, util in gpus for _ in range(3)] + lines.append(PROBE_APPS_MARKER) + lines += apps or [] + lines += [f"{PROBE_APPS_RC_PREFIX}0", PROBE_OWNERS_MARKER, PROBE_END_MARKER] + return ("\n".join(lines) + "\n").encode() + + +class FakeSsh: + """`_run_ssh` の代わり．接続先アドレスごとに応答を返し，呼び出しを記録する.""" + + def __init__(self, responses: dict[str, Callable[[list[str]], cli._SshResult]]) -> None: + self.responses = responses + self.calls: list[list[str]] = [] + self._lock = threading.Lock() + + def __call__( + self, argv: list[str], *, input_bytes: bytes | None = None, timeout: float | None = None + ) -> cli._SshResult: + with self._lock: + self.calls.append(list(argv)) + for address, respond in self.responses.items(): + if f"hasegawa@{address}" in argv: + return respond(argv) + raise AssertionError("想定外の接続先") + + +def _ok(stdout: bytes) -> Callable[[list[str]], cli._SshResult]: + return lambda argv: cli._SshResult(0, stdout, b"") + + +def _fail(rc: int = 255, stderr: bytes = b"ssh: connect to host: Connection timed out"): + return lambda argv: cli._SshResult(rc, b"", stderr) + + +@pytest.fixture +def project(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + """手元のプロジェクトルートを模したディレクトリ（ssh は一切起動させない）.""" + root = tmp_path / "proj" + (root / "src" / "tias_judge_model").mkdir(parents=True) + (root / "src" / "tias_judge_model" / "__init__.py").write_bytes(b"") + (root / "pyproject.toml").write_bytes(b"[project]\nname = 'x'\n") + (root / "configs").mkdir() + (root / "configs" / "remote.yaml").write_text(CONFIG_YAML, encoding="utf-8") + monkeypatch.chdir(root) + for var in ("TIAS_REMOTE_JUMP", "TIAS_REMOTE_ALLOW_SELF"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("TIAS_REMOTE_SSH", "ssh") + monkeypatch.setattr(cli.socket, "gethostname", lambda: "my-laptop") + + def _no_process(*args: object, **kwargs: object) -> None: + raise AssertionError("テストが実プロセス（ssh）を起動しようとした") + + monkeypatch.setattr(cli, "_popen", _no_process) + monkeypatch.setattr(cli.subprocess, "Popen", _no_process) + return root + + +def _spy_selection(monkeypatch: pytest.MonkeyPatch) -> list: + picked: list = [] + original = cli.select_gpu + + def spy(*args, **kwargs): + result = original(*args, **kwargs) + picked.append(result) + return result + + monkeypatch.setattr(cli, "select_gpu", spy) + return picked + + +# --------------------------------------------------------------------------- +# 接続できないサーバ +# --------------------------------------------------------------------------- + + +class TestUnreachableServers: + def test_gpus_continues_with_other_server( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + fake = FakeSsh({NEWTON_ADDR: _fail(), ARCHIMEDES_ADDR: _ok(_probe_stdout([(0, 5, 0)]))}) + monkeypatch.setattr(cli, "_run_ssh", fake) + picked = _spy_selection(monkeypatch) + assert cli.main(["gpus"]) == 0 + assert picked[0] is not None + assert (picked[0].server.name, picked[0].gpu_index) == ("Archimedes", 0) + assert len(fake.calls) == 2 + + def test_run_selects_gpu_on_reachable_server( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + fake = FakeSsh({NEWTON_ADDR: _fail(), ARCHIMEDES_ADDR: _ok(_probe_stdout([(0, 5, 0)]))}) + monkeypatch.setattr(cli, "_run_ssh", fake) + picked = _spy_selection(monkeypatch) + rc = cli.main(["run", "--name", "exp1", "--dry-run", "--", "python", "-V"]) + assert rc == 0 + assert picked[0].server.name == "Archimedes" + assert picked[0].state is GpuState.UNUSED + + def test_unparsable_probe_is_treated_as_unreachable( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + fake = FakeSsh( + {NEWTON_ADDR: _ok(b"garbage\n"), ARCHIMEDES_ADDR: _ok(_probe_stdout([(1, 5, 0)]))} + ) + monkeypatch.setattr(cli, "_run_ssh", fake) + picked = _spy_selection(monkeypatch) + assert cli.main(["gpus"]) == 0 + assert (picked[0].server.name, picked[0].gpu_index) == ("Archimedes", 1) + + @pytest.mark.parametrize("subcommand", [["gpus"], ["status"], ["push"]]) + def test_all_unreachable_is_error( + self, project: Path, monkeypatch: pytest.MonkeyPatch, subcommand: list[str] + ) -> None: + fake = FakeSsh({NEWTON_ADDR: _fail(), ARCHIMEDES_ADDR: _fail()}) + monkeypatch.setattr(cli, "_run_ssh", fake) + assert cli.main(subcommand) != 0 + + def test_run_all_unreachable_is_error( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + fake = FakeSsh({NEWTON_ADDR: _fail(), ARCHIMEDES_ADDR: _fail()}) + monkeypatch.setattr(cli, "_run_ssh", fake) + assert cli.main(["run", "--name", "exp1", "--dry-run", "--", "python"]) != 0 + + +class TestRunSelection: + def test_unused_on_lower_priority_server_is_chosen( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + newton = _probe_stdout([(0, 300, 0)], apps=["GPU-0, 4321, 300"]) + fake = FakeSsh({NEWTON_ADDR: _ok(newton), ARCHIMEDES_ADDR: _ok(_probe_stdout([(0, 5, 0)]))}) + monkeypatch.setattr(cli, "_run_ssh", fake) + picked = _spy_selection(monkeypatch) + assert cli.main(["run", "--name", "exp1", "--dry-run", "--", "python"]) == 0 + assert picked[0].server.name == "Archimedes" + + def test_no_free_gpu_does_not_start_job( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + busy = _probe_stdout([(0, 20000, 95)]) + fake = FakeSsh({NEWTON_ADDR: _ok(busy), ARCHIMEDES_ADDR: _ok(busy)}) + monkeypatch.setattr(cli, "_run_ssh", fake) + assert cli.main(["run", "--name", "exp1", "--", "python"]) != 0 + # probe の 2 回だけで，起動のための ssh は呼ばない + assert len(fake.calls) == 2 + + def test_gpu_without_server_is_rejected_before_ssh( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + fake = FakeSsh({}) + monkeypatch.setattr(cli, "_run_ssh", fake) + assert cli.main(["run", "--name", "exp1", "--gpu", "0", "--", "python"]) != 0 + assert fake.calls == [] + + def test_explicit_busy_gpu_is_error_without_force( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + busy = _probe_stdout([(0, 20000, 95)]) + fake = FakeSsh({NEWTON_ADDR: _ok(busy), ARCHIMEDES_ADDR: _ok(busy)}) + monkeypatch.setattr(cli, "_run_ssh", fake) + args = ["run", "--name", "e", "--server", "Newton", "--gpu", "0", "--dry-run", "--", "x"] + assert cli.main(args) != 0 + assert cli.main([*args[:7], "--force", *args[7:]]) == 0 + + +# --------------------------------------------------------------------------- +# 環境変数 +# --------------------------------------------------------------------------- + + +class TestEnvironment: + def _run_gpus(self, monkeypatch: pytest.MonkeyPatch) -> FakeSsh: + fake = FakeSsh( + { + NEWTON_ADDR: _ok(_probe_stdout([(0, 5, 0)])), + ARCHIMEDES_ADDR: _ok(_probe_stdout([(0, 5, 0)])), + } + ) + monkeypatch.setattr(cli, "_run_ssh", fake) + assert cli.main(["gpus"]) == 0 + return fake + + def test_config_jump_host_is_used(self, project: Path, monkeypatch: pytest.MonkeyPatch) -> None: + fake = self._run_gpus(monkeypatch) + for argv in fake.calls: + assert argv[argv.index("-J") + 1] == "labpc" + assert "BatchMode=yes" in argv + + def test_empty_jump_env_means_direct( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("TIAS_REMOTE_JUMP", "") + fake = self._run_gpus(monkeypatch) + assert fake.calls + for argv in fake.calls: + assert "-J" not in argv + + def test_self_server_is_detected(self, project: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(cli.socket, "gethostname", lambda: "Newton") + fake = FakeSsh({}) + monkeypatch.setattr(cli, "_run_ssh", fake) + assert cli.main(["gpus"]) != 0 + assert fake.calls == [] + + def test_allow_self_disables_detection( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(cli.socket, "gethostname", lambda: "Newton") + monkeypatch.setenv("TIAS_REMOTE_ALLOW_SELF", "1") + fake = self._run_gpus(monkeypatch) + assert len(fake.calls) == 2 + + +# --------------------------------------------------------------------------- +# pull +# --------------------------------------------------------------------------- + +LISTING_NAMES = [ + "config.yaml", + "history.csv", + "metrics.json", + "confusion_matrix.csv", + "log.txt", + "figures/confusion_matrix.png", + "split.json", + "predictions.csv", + "P0001_yamada.png", + "manifest.csv", + "checkpoints/best.pt", +] +ALLOWED = LISTING_NAMES[:6] +SECRET_NAMES = ["split.json", "predictions.csv", "P0001_yamada", "manifest.csv"] + + +def _listing() -> bytes: + return b"".join(f"f {RUN_ID}/{name}\0".encode() for name in LISTING_NAMES) + + +def _tar_of(names: list[str]) -> bytes: + buffer = io.BytesIO() + with tarfile.open(fileobj=buffer, mode="w", format=tarfile.PAX_FORMAT) as tar: + for name in names: + data = f"content of {name}".encode() + info = tarfile.TarInfo(name) + info.size = len(data) + tar.addfile(info, io.BytesIO(data)) + return buffer.getvalue() + + +class FakeReceive: + """`_receive_tar` の代わり．要求されたパスの tar（と任意の追加メンバー）を返す.""" + + def __init__(self, extra: list[str] | None = None) -> None: + self.extra = extra or [] + self.requested: list[str] = [] + + def __call__(self, argv: list[str], payload: bytes, destination) -> cli._SshResult: + self.requested = [p.decode() for p in payload.split(b"\0") if p] + destination.write(_tar_of(self.requested + self.extra)) + return cli._SshResult(0, b"", b"") + + +def _setup_pull(monkeypatch: pytest.MonkeyPatch, receive: FakeReceive) -> FakeSsh: + fake = FakeSsh( + { + NEWTON_ADDR: _ok(_listing()), + ARCHIMEDES_ADDR: _fail(rc=EXIT_NOT_FOUND, stderr=b""), + } + ) + monkeypatch.setattr(cli, "_run_ssh", fake) + monkeypatch.setattr(cli, "_receive_tar", receive) + return fake + + +def _local_files(root: Path) -> list[str]: + run = root / "runs" / RUN_ID + if not run.exists(): + return [] + return sorted(p.relative_to(run).as_posix() for p in run.rglob("*") if p.is_file()) + + +class TestPull: + def test_pulls_only_whitelisted_files( + self, + project: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + caplog: pytest.LogCaptureFixture, + ) -> None: + caplog.set_level(logging.DEBUG) + receive = FakeReceive() + _setup_pull(monkeypatch, receive) + assert cli.main(["pull", RUN_ID]) == 0 + assert sorted(receive.requested) == sorted(f"{RUN_ID}/{n}" for n in ALLOWED) + assert _local_files(project) == sorted(ALLOWED) + captured = capsys.readouterr() + output = captured.out + captured.err + caplog.text + for secret in [*SECRET_NAMES, "best.pt"]: + assert secret not in output + # 回収しなかった件数（split・predictions・画像・manifest・checkpoint の 5 件）は示す + assert "5 件" in output + + def test_with_checkpoints(self, project: Path, monkeypatch: pytest.MonkeyPatch) -> None: + receive = FakeReceive() + _setup_pull(monkeypatch, receive) + assert cli.main(["pull", RUN_ID, "--with-checkpoints"]) == 0 + assert _local_files(project) == sorted([*ALLOWED, "checkpoints/best.pt"]) + + @pytest.mark.parametrize( + "extra", + [ + f"{RUN_ID}/split.json", + f"{RUN_ID}/P0001_yamada.png", + f"{RUN_ID}/../evil.txt", + "../evil.txt", + "other_run/config.yaml", + ], + ) + def test_malicious_archive_extracts_nothing( + self, + project: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + caplog: pytest.LogCaptureFixture, + extra: str, + ) -> None: + caplog.set_level(logging.DEBUG) + _setup_pull(monkeypatch, FakeReceive(extra=[extra])) + assert cli.main(["pull", RUN_ID]) != 0 + assert _local_files(project) == [] + assert not (project / "runs" / "evil.txt").exists() + assert not (project / "evil.txt").exists() + captured = capsys.readouterr() + output = captured.out + captured.err + caplog.text + assert "P0001_yamada" not in output + assert "split.json" not in output + + @pytest.mark.parametrize("run_id", ["../etc", "a/b", "a\\b", "..", "x..y"]) + def test_invalid_run_id_is_rejected_before_ssh( + self, project: Path, monkeypatch: pytest.MonkeyPatch, run_id: str + ) -> None: + fake = FakeSsh({}) + monkeypatch.setattr(cli, "_run_ssh", fake) + assert cli.main(["pull", run_id]) != 0 + assert fake.calls == [] + + def test_run_on_multiple_servers_is_error( + self, project: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + fake = FakeSsh({NEWTON_ADDR: _ok(_listing()), ARCHIMEDES_ADDR: _ok(_listing())}) + monkeypatch.setattr(cli, "_run_ssh", fake) + monkeypatch.setattr(cli, "_receive_tar", FakeReceive()) + assert cli.main(["pull", RUN_ID]) != 0 + assert cli.main(["pull", RUN_ID, "--server", "Newton"]) == 0 + + def test_run_not_found_is_error(self, project: Path, monkeypatch: pytest.MonkeyPatch) -> None: + missing = _fail(rc=EXIT_NOT_FOUND, stderr=b"") + fake = FakeSsh({NEWTON_ADDR: missing, ARCHIMEDES_ADDR: missing}) + monkeypatch.setattr(cli, "_run_ssh", fake) + assert cli.main(["pull", RUN_ID]) != 0 + + +# --------------------------------------------------------------------------- +# push +# --------------------------------------------------------------------------- + + +class TestPush: + def test_dry_run_lists_code_but_not_excluded( + self, project: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + for rel in ["data/synthetic/manifest.csv", "runs/x/metrics.json", ".claude/a.md"]: + path = project / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"x") + fake = FakeSsh({}) + monkeypatch.setattr(cli, "_run_ssh", fake) + assert cli.main(["push", "--dry-run"]) == 0 + out = capsys.readouterr().out + assert "pyproject.toml" in out + assert "configs/remote.yaml" in out + assert "data/synthetic" not in out + assert "runs/x" not in out + assert ".claude" not in out + assert fake.calls == [] diff --git a/tests/remote/probe_helpers.py b/tests/remote/probe_helpers.py new file mode 100644 index 0000000..82d1e4a --- /dev/null +++ b/tests/remote/probe_helpers.py @@ -0,0 +1,74 @@ +"""テスト用に probe コマンド（SPEC_03「取得」）の出力文字列を組み立てる補助関数. + +実サーバの nvidia-smi・ps を呼ばず，出力の文字列だけで判定を固定するために使う． +""" + +from __future__ import annotations + +from collections.abc import Sequence + +from tias_judge_model.remote.commands import ( + PROBE_APPS_MARKER, + PROBE_APPS_RC_PREFIX, + PROBE_END_MARKER, + PROBE_GPU_RC_PREFIX, + PROBE_GPUS_MARKER, + PROBE_JOB_PREFIX, + PROBE_OTHER_PREFIX, + PROBE_OWNERS_MARKER, + PROBE_PMON_MARKER, + PROBE_SAMPLES_MARKER, + PROBE_SESSIONS_MARKER, +) + +PMON_HEADER = [ + "# gpu pid type sm mem enc dec jpg ofa command", + "# Idx # C/G % % % % % % name", +] + + +def gpu_line(index: int, uuid: str, used: int, util: int, *, total: int = 49140) -> str: + """`nvidia-smi --query-gpu=index,uuid,name,memory.used,memory.total,utilization.gpu` の 1 行.""" + return f"{index}, {uuid}, NVIDIA RTX A6000, {used}, {total}, {util}" + + +def sample_lines(index: int, utils: Sequence[int], memory: int = 10) -> list[str]: + """観測サンプル（`index, 使用率, 使用中メモリ`）の行.""" + return [f"{index}, {util}, {memory}" for util in utils] + + +def probe_text( + gpus: Sequence[str], + *, + gpu_rc: int = 0, + samples: Sequence[str] | None = None, + pmon: Sequence[str] | None = None, + apps: Sequence[str] | None = (), + apps_rc: int = 0, + owners: Sequence[str] = (), + jobs: Sequence[str] = (), + others: Sequence[str] = (), + preamble: Sequence[str] = (), +) -> str: + """probe の出力を組み立てる（samples・pmon が None なら節を出さない．apps が None なら + 計算プロセスの節ごと出さない）.""" + lines = [*preamble, PROBE_SESSIONS_MARKER] + lines += [f"{PROBE_JOB_PREFIX}{name}" for name in jobs] + lines += [f"{PROBE_OTHER_PREFIX}{name}" for name in others] + lines.append(PROBE_GPUS_MARKER) + lines += list(gpus) + lines.append(f"{PROBE_GPU_RC_PREFIX}{gpu_rc}") + if samples is not None: + lines.append(PROBE_SAMPLES_MARKER) + lines += list(samples) + if pmon is not None: + lines.append(PROBE_PMON_MARKER) + lines += list(pmon) + if apps is not None: + lines.append(PROBE_APPS_MARKER) + lines += list(apps) + lines.append(f"{PROBE_APPS_RC_PREFIX}{apps_rc}") + lines.append(PROBE_OWNERS_MARKER) + lines += list(owners) + lines.append(PROBE_END_MARKER) + return "\n".join(lines) + "\n" diff --git a/tests/remote/test_commands.py b/tests/remote/test_commands.py new file mode 100644 index 0000000..58c2d47 --- /dev/null +++ b/tests/remote/test_commands.py @@ -0,0 +1,114 @@ +"""ssh 引数とサーバ上で実行するコマンドの組み立てのテスト（SPEC_03「接続」「取得」「run」）.""" + +from __future__ import annotations + +import pytest + +from tias_judge_model.remote.commands import ( + build_job_env, + build_probe_command, + build_ssh_argv, + validate_job_name, +) + + +def _argv(jump: str | None, timeout: int = 15) -> list[str]: + return build_ssh_argv( + "ssh", + user="hasegawa", + address="192.168.23.14", + jump_host=jump, + connect_timeout_sec=timeout, + remote_command="echo hi", + ) + + +def _has_option(argv: list[str], value: str) -> bool: + return any(argv[i] == "-o" and argv[i + 1] == value for i in range(len(argv) - 1)) + + +class TestBuildSshArgv: + def test_with_jump_host(self) -> None: + argv = _argv("labpc") + assert argv[0] == "ssh" + j = argv.index("-J") + assert argv[j + 1] == "labpc" + assert "hasegawa@192.168.23.14" in argv + assert argv[-1] == "echo hi" + # 接続先より前にオプションがある + assert j < argv.index("hasegawa@192.168.23.14") + + def test_without_jump_host(self) -> None: + argv = _argv(None) + assert "-J" not in argv + assert "hasegawa@192.168.23.14" in argv + + @pytest.mark.parametrize("jump", ["labpc", None]) + def test_batch_mode_and_timeout_always(self, jump: str | None) -> None: + argv = _argv(jump, timeout=7) + assert _has_option(argv, "BatchMode=yes") + assert _has_option(argv, "ConnectTimeout=7") + + def test_remote_command_is_single_argument(self) -> None: + argv = build_ssh_argv( + "ssh", + user="u", + address="h", + jump_host=None, + connect_timeout_sec=5, + remote_command="cd ~/x; echo a b", + ) + assert argv[-1] == "cd ~/x; echo a b" + + @pytest.mark.parametrize("bad", ["-oProxyCommand=evil", "a b", ""]) + def test_rejects_option_like_destination(self, bad: str) -> None: + with pytest.raises(ValueError): + build_ssh_argv( + "ssh", + user="u", + address=bad, + jump_host=None, + connect_timeout_sec=5, + remote_command="true", + ) + + +class TestBuildProbeCommand: + def test_queries_listed_in_spec(self) -> None: + command = build_probe_command("~/work/x", observe_sec=10, sample_interval_ms=500) + assert "--query-gpu=index,uuid,name,memory.used,memory.total,utilization.gpu" in command + assert "--query-compute-apps=gpu_uuid,pid,used_memory" in command + assert "--format=csv,noheader,nounits" in command + + def test_observation_uses_interval_and_duration(self) -> None: + command = build_probe_command("~/work/x", observe_sec=10, sample_interval_ms=500) + assert "-lms 500" in command + assert "timeout 10" in command + assert "pmon -s u" in command + + def test_observe_zero_does_not_observe(self) -> None: + command = build_probe_command("~/work/x", observe_sec=0, sample_interval_ms=500) + assert "-lms" not in command + assert "pmon" not in command + assert "--query-gpu=index,uuid,name" in command + + def test_is_single_line(self) -> None: + command = build_probe_command("~/work/x", observe_sec=10, sample_interval_ms=500) + assert "\n" not in command + + +class TestJobSettings: + @pytest.mark.parametrize("name", ["a", "exp1", "Exp_1-b", "x" * 64]) + def test_valid_job_names(self, name: str) -> None: + assert validate_job_name(name) == name + + @pytest.mark.parametrize("name", ["", "_a", "-a", "a b", "a.b", "a:b", "x" * 65, "a/b"]) + def test_invalid_job_names(self, name: str) -> None: + with pytest.raises(ValueError): + validate_job_name(name) + + def test_job_env_contains_selected_gpu_and_extra(self) -> None: + env = build_job_env(2, ["FOO=bar", "EMPTY="]) + assert env["CUDA_VISIBLE_DEVICES"] == "2" + assert env["FOO"] == "bar" + assert env["EMPTY"] == "" diff --git a/tests/remote/test_config.py b/tests/remote/test_config.py new file mode 100644 index 0000000..e355c60 --- /dev/null +++ b/tests/remote/test_config.py @@ -0,0 +1,228 @@ +"""サーバ設定の検証と環境変数の上書きのテスト（SPEC_03「サーバ設定」「サーバ上での誤実行の防止」）.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import yaml +from pydantic import ValidationError + +from tias_judge_model.remote.config import ( + detect_self_server, + is_self_check_disabled, + load_remote_config, + resolve_jump_host, +) + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _base() -> dict: + return { + "jump_host": "labpc", + "connect_timeout_sec": 15, + "gpu_free": { + "max_memory_used_mib": 1024, + "max_utilization_percent": 10, + "observe_sec": 10, + "sample_interval_ms": 500, + }, + "servers": [ + { + "name": "Newton", + "address": "192.168.23.14", + "user": "hasegawa", + "priority": 1, + "remote_dir": "~/work/TIASJudgeModel", + "data_dir": "~/data/TIASJudgeModel", + "uv": "~/.local/bin/uv", + "data_ready": True, + }, + { + "name": "Archimedes", + "address": "192.168.23.15", + "user": "hasegawa", + "priority": 2, + "remote_dir": "~/work/TIASJudgeModel", + "data_dir": "~/data/TIASJudgeModel", + "uv": "~/.local/bin/uv", + "data_ready": False, + }, + ], + } + + +def _write(tmp_path: Path, data: dict) -> Path: + path = tmp_path / "remote.yaml" + path.write_text(yaml.safe_dump(data, allow_unicode=True), encoding="utf-8") + return path + + +def test_repository_remote_yaml_loads() -> None: + config = load_remote_config(REPO_ROOT / "configs" / "remote.yaml") + assert len(config.servers) >= 1 + assert {s.name for s in config.servers} >= {"Newton"} + + +def test_valid_config_loads(tmp_path: Path) -> None: + config = load_remote_config(_write(tmp_path, _base())) + assert [s.name for s in config.servers] == ["Newton", "Archimedes"] + assert config.jump_host == "labpc" + assert config.gpu_free.observe_sec == 10 + assert config.servers[1].data_ready is False + + +def test_jump_host_null_means_direct(tmp_path: Path) -> None: + data = _base() + data["jump_host"] = None + assert load_remote_config(_write(tmp_path, data)).jump_host is None + + +def _invalid(tmp_path: Path, data: dict) -> None: + with pytest.raises(ValidationError): + load_remote_config(_write(tmp_path, data)) + + +def test_servers_must_not_be_empty(tmp_path: Path) -> None: + data = _base() + data["servers"] = [] + _invalid(tmp_path, data) + + +def test_servers_are_required(tmp_path: Path) -> None: + data = _base() + del data["servers"] + _invalid(tmp_path, data) + + +def test_duplicate_server_names_are_rejected(tmp_path: Path) -> None: + data = _base() + data["servers"][1]["name"] = "Newton" + _invalid(tmp_path, data) + + +@pytest.mark.parametrize("name", ["New ton", "Newton;rm", "サーバ", "a/b", ""]) +def test_server_name_characters(tmp_path: Path, name: str) -> None: + data = _base() + data["servers"][0]["name"] = name + _invalid(tmp_path, data) + + +@pytest.mark.parametrize("name", ["Newton_2", "gpu-01", "A1"]) +def test_server_name_allowed_characters(tmp_path: Path, name: str) -> None: + data = _base() + data["servers"][0]["name"] = name + assert load_remote_config(_write(tmp_path, data)).servers[0].name == name + + +@pytest.mark.parametrize("priority", ["high", 1.5]) +def test_priority_must_be_integer(tmp_path: Path, priority: object) -> None: + data = _base() + data["servers"][0]["priority"] = priority + _invalid(tmp_path, data) + + +@pytest.mark.parametrize( + ("key", "value", "ok"), + [ + ("max_memory_used_mib", 0, True), + ("max_memory_used_mib", -1, False), + ("max_utilization_percent", 0, True), + ("max_utilization_percent", 100, True), + ("max_utilization_percent", -1, False), + ("max_utilization_percent", 101, False), + ("observe_sec", 0, True), + ("observe_sec", 120, True), + ("observe_sec", -1, False), + ("observe_sec", 121, False), + ("sample_interval_ms", 100, True), + ("sample_interval_ms", 5000, True), + ("sample_interval_ms", 99, False), + ("sample_interval_ms", 5001, False), + ], +) +def test_gpu_free_ranges(tmp_path: Path, key: str, value: int, ok: bool) -> None: + data = _base() + data["gpu_free"][key] = value + if ok: + assert getattr(load_remote_config(_write(tmp_path, data)).gpu_free, key) == value + else: + _invalid(tmp_path, data) + + +@pytest.mark.parametrize(("value", "ok"), [(1, True), (0, False), (-5, False)]) +def test_connect_timeout_at_least_one(tmp_path: Path, value: int, ok: bool) -> None: + data = _base() + data["connect_timeout_sec"] = value + if ok: + assert load_remote_config(_write(tmp_path, data)).connect_timeout_sec == value + else: + _invalid(tmp_path, data) + + +def test_unknown_top_level_key_is_rejected(tmp_path: Path) -> None: + data = _base() + data["jump_hots"] = "labpc" + _invalid(tmp_path, data) + + +def test_unknown_server_key_is_rejected(tmp_path: Path) -> None: + data = _base() + data["servers"][0]["data_redy"] = True + _invalid(tmp_path, data) + + +def test_unknown_gpu_free_key_is_rejected(tmp_path: Path) -> None: + data = _base() + data["gpu_free"]["observe_seconds"] = 5 + _invalid(tmp_path, data) + + +def test_missing_required_server_field_is_rejected(tmp_path: Path) -> None: + data = _base() + del data["servers"][0]["address"] + _invalid(tmp_path, data) + + +def test_missing_file_is_value_error(tmp_path: Path) -> None: + with pytest.raises(ValueError): + load_remote_config(tmp_path / "missing.yaml") + + +class TestJumpHostEnv: + def _config(self, tmp_path: Path, jump: str | None = "labpc"): + data = _base() + data["jump_host"] = jump + return load_remote_config(_write(tmp_path, data)) + + def test_no_env_uses_config(self, tmp_path: Path) -> None: + assert resolve_jump_host(self._config(tmp_path), {}) == "labpc" + + def test_env_overrides(self, tmp_path: Path) -> None: + env = {"TIAS_REMOTE_JUMP": "otherjump"} + assert resolve_jump_host(self._config(tmp_path), env) == "otherjump" + + def test_empty_env_means_direct(self, tmp_path: Path) -> None: + assert resolve_jump_host(self._config(tmp_path), {"TIAS_REMOTE_JUMP": ""}) is None + + def test_env_sets_jump_when_config_is_null(self, tmp_path: Path) -> None: + config = self._config(tmp_path, None) + assert resolve_jump_host(config, {}) is None + assert resolve_jump_host(config, {"TIAS_REMOTE_JUMP": "labpc"}) == "labpc" + + +class TestSelfServer: + def test_hostname_matching_server_name_is_detected(self, tmp_path: Path) -> None: + config = load_remote_config(_write(tmp_path, _base())) + server = detect_self_server(config, "Newton") + assert server is not None and server.name == "Newton" + + def test_other_hostname_is_not_detected(self, tmp_path: Path) -> None: + config = load_remote_config(_write(tmp_path, _base())) + assert detect_self_server(config, "my-laptop") is None + + def test_allow_self_env(self) -> None: + assert is_self_check_disabled({"TIAS_REMOTE_ALLOW_SELF": "1"}) + assert not is_self_check_disabled({}) + assert not is_self_check_disabled({"TIAS_REMOTE_ALLOW_SELF": "0"}) diff --git a/tests/remote/test_files.py b/tests/remote/test_files.py new file mode 100644 index 0000000..1bdbbba --- /dev/null +++ b/tests/remote/test_files.py @@ -0,0 +1,86 @@ +"""push で送るファイルの選別のテスト（SPEC_03「push」）. + +`.git/`・`.venv/`・`__pycache__/`・`.pytest_cache/`・`.ruff_cache/` はどこにあっても送らない． +`data/`・`runs/`・`.claude/` はプロジェクト直下のみ除外する（`src/tias_judge_model/data/`・ +`tests/data/` は送る）． +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tias_judge_model.remote.files import collect_push_entries, is_excluded + +SENT = [ + "pyproject.toml", + "uv.lock", + ".python-version", + "src/tias_judge_model/__init__.py", + "src/tias_judge_model/data/manifest.py", + "src/tias_judge_model/remote/gpu.py", + "tests/data/test_split.py", + "tests/remote/test_gpu.py", + "scripts/make_synthetic_dataset.py", + "configs/remote.yaml", +] + +NOT_SENT = [ + ".git/config", + ".git/objects/ab/cdef", + ".venv/Scripts/python.exe", + "__pycache__/x.pyc", + "src/tias_judge_model/__pycache__/a.cpython-312.pyc", + "tests/data/__pycache__/b.pyc", + ".pytest_cache/v/cache/nodeids", + ".ruff_cache/0.1/x", + "data/synthetic/manifest.csv", + "data/synthetic/images/pale_1/P0001.png", + "runs/20260927-120000_x/metrics.json", + "runs/_remote_logs/exp1.log", + ".claude/settings.json", + ".claude/rules/testing.md", +] + + +def _make_tree(root: Path, paths: list[str]) -> None: + for rel in paths: + path = root / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"x") + + +def test_collect_push_entries_excludes_listed_directories(tmp_path: Path) -> None: + _make_tree(tmp_path, SENT + NOT_SENT) + arcnames = [entry.arcname for entry in collect_push_entries(tmp_path)] + assert sorted(arcnames) == sorted(SENT) + + +def test_arcnames_are_posix_relative(tmp_path: Path) -> None: + _make_tree(tmp_path, ["src/tias_judge_model/data/manifest.py"]) + entries = collect_push_entries(tmp_path) + assert [e.arcname for e in entries] == ["src/tias_judge_model/data/manifest.py"] + assert entries[0].size == 1 + + +@pytest.mark.parametrize("rel", NOT_SENT) +def test_is_excluded_for_not_sent(rel: str) -> None: + assert is_excluded(rel) + + +@pytest.mark.parametrize("rel", SENT) +def test_is_not_excluded_for_sent(rel: str) -> None: + assert not is_excluded(rel) + + +@pytest.mark.parametrize("rel", ["data", "runs", ".claude", ".git", ".venv"]) +def test_root_directories_themselves_are_excluded(rel: str) -> None: + assert is_excluded(rel) + + +@pytest.mark.parametrize( + "rel", ["src/tias_judge_model/data", "tests/data", "src/tias_judge_model/runs_util.py"] +) +def test_nested_data_is_not_excluded(rel: str) -> None: + assert not is_excluded(rel) diff --git a/tests/remote/test_gpu.py b/tests/remote/test_gpu.py new file mode 100644 index 0000000..bf80327 --- /dev/null +++ b/tests/remote/test_gpu.py @@ -0,0 +1,462 @@ +"""GPU の空き判定（4 段階）と選択のテスト（SPEC_03「GPU の空き判定と選択」）. + +- 観測中の最大値で閾値を判定する（瞬間値が 0% でも観測中に超えたら使用中） +- 計算プロセスを所有者を問わず数える．一覧が取れなければ閾値内の GPU は「空き（利用者あり）」 +- GPU の行が解釈できない・サンプル 0 件なら「状態不明」 +- 選択では「未使用」をサーバの優先順位より優先する +""" + +from __future__ import annotations + +import pytest +from probe_helpers import PMON_HEADER, gpu_line, probe_text, sample_lines + +from tias_judge_model.remote.config import GpuFreeThreshold, ServerConfig +from tias_judge_model.remote.gpu import ( + GpuInfo, + GpuSelection, + GpuState, + ProbeOutput, + ServerProbe, + check_selection_options, + classify_gpu, + gpu_users, + parse_probe_output, + select_gpu, +) +from tias_judge_model.remote.observe import UsagePattern + +OBSERVE = GpuFreeThreshold( + max_memory_used_mib=1024, max_utilization_percent=10, observe_sec=10, sample_interval_ms=500 +) +INSTANT = GpuFreeThreshold( + max_memory_used_mib=1024, max_utilization_percent=10, observe_sec=0, sample_interval_ms=500 +) + +UUID0 = "GPU-00000000-aaaa-bbbb-cccc-000000000000" +UUID1 = "GPU-11111111-aaaa-bbbb-cccc-111111111111" + + +def _states(output: ProbeOutput) -> dict[int | None, GpuState]: + return {gpu.index: gpu.state for gpu in output.gpus} + + +def _gpu(output: ProbeOutput, index: int) -> GpuInfo: + return next(g for g in output.gpus if g.index == index) + + +# --------------------------------------------------------------------------- +# 判定（classify_gpu） +# --------------------------------------------------------------------------- + + +class TestClassifyGpu: + def test_within_threshold_without_processes_is_unused(self) -> None: + assert classify_gpu(0, 0, OBSERVE, 0) is GpuState.UNUSED + + def test_within_threshold_with_processes_is_shared(self) -> None: + assert classify_gpu(500, 5, OBSERVE, 1) is GpuState.SHARED + + def test_unknown_process_list_is_shared_not_unused(self) -> None: + assert classify_gpu(0, 0, OBSERVE, None) is GpuState.SHARED + + @pytest.mark.parametrize( + ("memory", "util", "expected"), + [ + (1023, 9, GpuState.UNUSED), + (1024, 9, GpuState.BUSY), # 「未満」なのでちょうどは超えた扱い + (1023, 10, GpuState.BUSY), + (5000, 0, GpuState.BUSY), + (0, 100, GpuState.BUSY), + ], + ) + def test_threshold_boundaries(self, memory: int, util: int, expected: GpuState) -> None: + assert classify_gpu(memory, util, OBSERVE, 0) is expected + + def test_busy_even_if_process_list_unknown(self) -> None: + assert classify_gpu(0, 50, OBSERVE, None) is GpuState.BUSY + + +# --------------------------------------------------------------------------- +# probe 出力の解析と判定（parse_probe_output） +# --------------------------------------------------------------------------- + + +class TestObservedJudgement: + def test_max_during_observation_is_used_even_if_instant_is_zero(self) -> None: + # 瞬間値は 0% だが，観測中に 74% まで上がった（2026-09-27 Newton の実測を模す） + text = probe_text( + [gpu_line(0, UUID0, used=5, util=0)], + samples=sample_lines(0, [0, 0, 74, 0, 0, 30, 0, 0, 0, 0]), + ) + output = parse_probe_output(text, OBSERVE) + assert _states(output) == {0: GpuState.BUSY} + + def test_memory_peak_during_observation_makes_busy(self) -> None: + samples = ["0, 0, 10", "0, 0, 2000", "0, 0, 10"] + text = probe_text([gpu_line(0, UUID0, used=10, util=0)], samples=samples) + assert _states(parse_probe_output(text, OBSERVE)) == {0: GpuState.BUSY} + + def test_always_within_threshold_without_processes_is_unused(self) -> None: + text = probe_text( + [gpu_line(0, UUID0, used=5, util=0)], samples=sample_lines(0, [0, 3, 9, 0, 1]) + ) + assert _states(parse_probe_output(text, OBSERVE)) == {0: GpuState.UNUSED} + + def test_single_sample_over_threshold_is_busy(self) -> None: + text = probe_text( + [gpu_line(0, UUID0, used=5, util=0)], samples=sample_lines(0, [0, 0, 10, 0]) + ) + assert _states(parse_probe_output(text, OBSERVE)) == {0: GpuState.BUSY} + + def test_no_samples_for_gpu_is_unknown(self) -> None: + text = probe_text( + [gpu_line(0, UUID0, used=5, util=0), gpu_line(1, UUID1, used=5, util=0)], + samples=sample_lines(0, [0, 0, 0]), + ) + assert _states(parse_probe_output(text, OBSERVE)) == { + 0: GpuState.UNUSED, + 1: GpuState.UNKNOWN, + } + + def test_missing_samples_section_makes_all_unknown(self) -> None: + text = probe_text([gpu_line(0, UUID0, used=5, util=0)], samples=None) + assert _states(parse_probe_output(text, OBSERVE)) == {0: GpuState.UNKNOWN} + + def test_observe_zero_uses_instant_value(self) -> None: + text = probe_text([gpu_line(0, UUID0, used=5, util=0), gpu_line(1, UUID1, used=5, util=50)]) + output = parse_probe_output(text, INSTANT) + assert _states(output) == {0: GpuState.UNUSED, 1: GpuState.BUSY} + + def test_observe_zero_does_not_require_samples(self) -> None: + text = probe_text([gpu_line(0, UUID0, used=2000, util=0)]) + assert _states(parse_probe_output(text, INSTANT)) == {0: GpuState.BUSY} + + +class TestProcessesAndUnknown: + def test_processes_are_matched_by_gpu_uuid(self) -> None: + text = probe_text( + [gpu_line(0, UUID0, used=500, util=0), gpu_line(1, UUID1, used=5, util=0)], + samples=sample_lines(0, [0, 0]) + sample_lines(1, [0, 0]), + apps=[f"{UUID0}, 4321, 480"], + owners=["4321 alice 99720"], + ) + output = parse_probe_output(text, OBSERVE) + assert _states(output) == {0: GpuState.SHARED, 1: GpuState.UNUSED} + process = _gpu(output, 0).processes[0] + assert process.pid == 4321 + assert process.user == "alice" + assert process.elapsed_sec == 99720 + assert _gpu(output, 1).processes == () + + def test_own_process_also_counts(self) -> None: + # 自分の別ジョブが載っている GPU にも次のジョブを重ねない（所有者を問わず数える） + text = probe_text( + [gpu_line(0, UUID0, used=300, util=0)], + samples=sample_lines(0, [0, 0]), + apps=[f"{UUID0}, 100, 300"], + owners=["100 hasegawa 10"], + ) + output = parse_probe_output(text, OBSERVE) + gpu = _gpu(output, 0) + assert gpu.state is GpuState.SHARED + assert gpu_users(gpu, "hasegawa") == ["自分"] + assert gpu_users(gpu, "someone_else") == ["hasegawa"] + + def test_process_list_failure_makes_within_threshold_gpus_shared(self) -> None: + text = probe_text( + [ + gpu_line(0, UUID0, used=5, util=0), + gpu_line(1, UUID1, used=5, util=0), + ], + samples=sample_lines(0, [0, 0]) + sample_lines(1, [0, 90]), + apps=[], + apps_rc=9, + ) + output = parse_probe_output(text, OBSERVE) + assert not output.processes_query_ok + # 閾値内は未使用と誤認せず「空き（利用者あり）」，閾値を超えたものは使用中のまま + assert _states(output) == {0: GpuState.SHARED, 1: GpuState.BUSY} + + def test_unparsable_process_list_makes_gpus_shared(self) -> None: + text = probe_text( + [gpu_line(0, UUID0, used=5, util=0), gpu_line(1, UUID1, used=5, util=0)], + samples=sample_lines(0, [0]) + sample_lines(1, [0]), + apps=["this is not csv"], + ) + output = parse_probe_output(text, OBSERVE) + assert _states(output) == {0: GpuState.SHARED, 1: GpuState.SHARED} + + def test_missing_process_section_makes_gpus_shared(self) -> None: + text = probe_text( + [gpu_line(0, UUID0, used=5, util=0)], samples=sample_lines(0, [0]), apps=None + ) + assert _states(parse_probe_output(text, OBSERVE)) == {0: GpuState.SHARED} + + @pytest.mark.parametrize( + "line", + [ + f"0, {UUID0}, NVIDIA RTX A6000, [N/A], 49140, 0", + f"0, {UUID0}, NVIDIA RTX A6000, 5, 49140, [N/A]", + "0, garbage", + "garbage", + ], + ) + def test_unparsable_gpu_line_is_unknown(self, line: str) -> None: + text = probe_text( + [line, gpu_line(1, UUID1, used=5, util=0)], + samples=sample_lines(0, [0]) + sample_lines(1, [0]), + ) + output = parse_probe_output(text, OBSERVE) + assert output.gpus[0].state is GpuState.UNKNOWN + assert _gpu(output, 1).state is GpuState.UNUSED + + def test_gpu_query_failure_yields_no_candidates(self) -> None: + text = probe_text([], gpu_rc=127, apps=None) + output = parse_probe_output(text, OBSERVE) + assert not output.gpu_query_ok + assert output.gpus == () + + def test_graphics_process_is_not_counted_from_pmon(self) -> None: + text = probe_text( + [gpu_line(0, UUID0, used=300, util=5)], + samples=sample_lines(0, [5, 0]), + pmon=[ + *PMON_HEADER, + " 0 1234 C 74 10 - - - - python", + " 0 1234 C 3 1 - - - - python", + " 0 999 G 50 2 - - - - Xorg", + " 0 - - - - - - - - -", + ], + apps=[f"{UUID0}, 1234, 280"], + owners=["1234 alice 45"], + ) + output = parse_probe_output(text, OBSERVE) + gpu = _gpu(output, 0) + assert [p.pid for p in gpu.processes] == [1234] + assert gpu.processes[0].max_sm_percent == 74 + + def test_preamble_before_markers_is_ignored(self) -> None: + text = probe_text( + [gpu_line(0, UUID0, used=5, util=0)], + samples=sample_lines(0, [0]), + preamble=["Welcome to Ubuntu", "Last login: Sat Sep 27 10:00:00 2026"], + ) + assert _states(parse_probe_output(text, OBSERVE)) == {0: GpuState.UNUSED} + + def test_truncated_output_is_error(self) -> None: + text = probe_text([gpu_line(0, UUID0, used=5, util=0)], samples=sample_lines(0, [0])) + truncated = text.split("@@tias-end")[0] + with pytest.raises(ValueError): + parse_probe_output(truncated, OBSERVE) + + def test_jobs_and_other_sessions(self) -> None: + text = probe_text([], jobs=["exp1"], others=["mywork"], samples=[]) + output = parse_probe_output(text, OBSERVE) + assert output.jobs == ("exp1",) + assert output.other_sessions == ("mywork",) + + +class TestPatternDoesNotAffectState: + def test_intermittent_but_within_threshold_is_unused(self) -> None: + text = probe_text( + [gpu_line(0, UUID0, used=5, util=0)], samples=sample_lines(0, [0, 5, 0, 0, 0]) + ) + gpu = _gpu(parse_probe_output(text, OBSERVE), 0) + assert gpu.pattern is UsagePattern.INTERMITTENT + assert gpu.state is GpuState.UNUSED + + def test_constant_but_within_threshold_is_unused(self) -> None: + text = probe_text( + [gpu_line(0, UUID0, used=5, util=3)], samples=sample_lines(0, [3, 3, 3, 3, 3]) + ) + gpu = _gpu(parse_probe_output(text, OBSERVE), 0) + assert gpu.pattern is UsagePattern.CONSTANT + assert gpu.state is GpuState.UNUSED + + def test_memory_only_pattern_is_shared(self) -> None: + text = probe_text( + [gpu_line(0, UUID0, used=500, util=0)], + samples=sample_lines(0, [0, 0, 0], memory=500), + apps=[f"{UUID0}, 1, 500"], + owners=["1 bob 100"], + ) + gpu = _gpu(parse_probe_output(text, OBSERVE), 0) + assert gpu.pattern is UsagePattern.MEMORY_ONLY + assert gpu.state is GpuState.SHARED + + def test_idle_pattern_but_memory_over_threshold_is_busy(self) -> None: + text = probe_text( + [gpu_line(0, UUID0, used=2000, util=0)], + samples=sample_lines(0, [0, 0, 0], memory=2000), + ) + gpu = _gpu(parse_probe_output(text, OBSERVE), 0) + assert gpu.pattern is UsagePattern.IDLE + assert gpu.state is GpuState.BUSY + + +# --------------------------------------------------------------------------- +# 選択（select_gpu） +# --------------------------------------------------------------------------- + + +def _server(name: str, priority: int, *, data_ready: bool = True) -> ServerConfig: + return ServerConfig( + name=name, + address="192.168.0.1", + user="hasegawa", + priority=priority, + remote_dir="~/work/TIASJudgeModel", + data_dir="~/data/TIASJudgeModel", + data_ready=data_ready, + ) + + +def _info(index: int, state: GpuState) -> GpuInfo: + return GpuInfo( + index=index, + uuid=f"GPU-{index}", + name="GPU", + memory_used_mib=0, + memory_total_mib=49140, + utilization_percent=0, + state=state, + processes_known=True, + ) + + +def _probe(server: ServerConfig, *states: GpuState, gpu_query_ok: bool = True) -> ServerProbe: + gpus = tuple(_info(i, s) for i, s in enumerate(states)) + return ServerProbe( + server, + ProbeOutput( + jobs=(), + other_sessions=(), + gpu_query_ok=gpu_query_ok, + processes_query_ok=True, + gpus=gpus if gpu_query_ok else (), + ), + ) + + +def _unreachable(server: ServerConfig) -> ServerProbe: + return ServerProbe(server, None, "Connection timed out") + + +U, S, B, X = GpuState.UNUSED, GpuState.SHARED, GpuState.BUSY, GpuState.UNKNOWN +NEWTON = _server("Newton", 1) +ARCHIMEDES = _server("Archimedes", 2) + + +def _picked(selection: GpuSelection | None) -> tuple[str, int, GpuState] | None: + if selection is None: + return None + return selection.server.name, selection.gpu_index, selection.state + + +class TestSelectGpuAutomatic: + def test_unused_beats_server_priority(self) -> None: + # Newton（優先）は利用者ありのみ，Archimedes に未使用 → Archimedes を選ぶ + probes = [_probe(NEWTON, S, S, B), _probe(ARCHIMEDES, B, U)] + assert _picked(select_gpu(probes)) == ("Archimedes", 1, U) + + def test_priority_order_among_unused(self) -> None: + probes = [_probe(ARCHIMEDES, U), _probe(NEWTON, B, U)] + assert _picked(select_gpu(probes)) == ("Newton", 1, U) + + def test_same_priority_uses_config_order(self) -> None: + first = _server("Alpha", 1) + second = _server("Beta", 1) + assert _picked(select_gpu([_probe(first, U), _probe(second, U)])) == ("Alpha", 0, U) + assert _picked(select_gpu([_probe(second, U), _probe(first, U)])) == ("Beta", 0, U) + + def test_lowest_index_within_server(self) -> None: + probe = _probe(NEWTON, B, U, U) + assert _picked(select_gpu([probe])) == ("Newton", 1, U) + + def test_index_order_even_if_listed_out_of_order(self) -> None: + output = ProbeOutput( + jobs=(), + other_sessions=(), + gpu_query_ok=True, + processes_query_ok=True, + gpus=(_info(3, U), _info(1, U), _info(2, U)), + ) + assert _picked(select_gpu([ServerProbe(NEWTON, output)])) == ("Newton", 1, U) + + def test_shared_when_no_unused_by_priority(self) -> None: + probes = [_probe(NEWTON, B, S), _probe(ARCHIMEDES, S)] + assert _picked(select_gpu(probes)) == ("Newton", 1, S) + + def test_none_when_only_busy_or_unknown(self) -> None: + probes = [_probe(NEWTON, B, X), _probe(ARCHIMEDES, X, B)] + assert select_gpu(probes) is None + + def test_data_not_ready_excluded_by_default(self) -> None: + archimedes = _server("Archimedes", 2, data_ready=False) + probes = [_probe(NEWTON, S), _probe(archimedes, U)] + assert _picked(select_gpu(probes)) == ("Newton", 0, S) + assert _picked(select_gpu(probes, include_not_ready=True)) == ("Archimedes", 0, U) + + def test_only_data_not_ready_has_free_gpu(self) -> None: + archimedes = _server("Archimedes", 2, data_ready=False) + probes = [_probe(NEWTON, B), _probe(archimedes, U)] + assert select_gpu(probes) is None + assert _picked(select_gpu(probes, include_not_ready=True)) == ("Archimedes", 0, U) + + def test_server_option_restricts_candidates(self) -> None: + probes = [_probe(NEWTON, U), _probe(ARCHIMEDES, B, S)] + assert _picked(select_gpu(probes, server_name="Archimedes")) == ("Archimedes", 1, S) + + def test_server_option_none_available(self) -> None: + probes = [_probe(NEWTON, U), _probe(ARCHIMEDES, B)] + assert select_gpu(probes, server_name="Archimedes") is None + + def test_unreachable_server_is_skipped(self) -> None: + probes = [_unreachable(NEWTON), _probe(ARCHIMEDES, S)] + assert _picked(select_gpu(probes)) == ("Archimedes", 0, S) + + def test_server_with_failed_gpu_query_is_skipped(self) -> None: + probes = [_probe(NEWTON, gpu_query_ok=False), _probe(ARCHIMEDES, S)] + assert _picked(select_gpu(probes)) == ("Archimedes", 0, S) + + def test_unknown_server_name_is_error(self) -> None: + with pytest.raises(ValueError): + select_gpu([_probe(NEWTON, U)], server_name="Nowhere") + + +class TestSelectGpuExplicit: + @pytest.mark.parametrize("state", [U, S]) + def test_unused_or_shared_is_accepted(self, state: GpuState) -> None: + probes = [_probe(NEWTON, B, state)] + selection = select_gpu(probes, server_name="Newton", gpu_index=1) + assert _picked(selection) == ("Newton", 1, state) + assert selection is not None and not selection.forced + + @pytest.mark.parametrize("state", [B, X]) + def test_busy_or_unknown_is_error_without_force(self, state: GpuState) -> None: + probes = [_probe(NEWTON, U, state)] + with pytest.raises(ValueError): + select_gpu(probes, server_name="Newton", gpu_index=1) + + @pytest.mark.parametrize("state", [B, X]) + def test_busy_or_unknown_is_allowed_with_force(self, state: GpuState) -> None: + probes = [_probe(NEWTON, U, state)] + selection = select_gpu(probes, server_name="Newton", gpu_index=1, force=True) + assert _picked(selection) == ("Newton", 1, state) + assert selection is not None and selection.forced + + def test_explicit_gpu_even_if_other_server_has_unused(self) -> None: + probes = [_probe(NEWTON, S), _probe(ARCHIMEDES, U)] + assert _picked(select_gpu(probes, server_name="Newton", gpu_index=0)) == ("Newton", 0, S) + + def test_gpu_without_server_is_error(self) -> None: + with pytest.raises(ValueError): + select_gpu([_probe(NEWTON, U)], gpu_index=0) + with pytest.raises(ValueError): + check_selection_options(server_name=None, gpu_index=0, force=False) + + def test_gpu_with_server_passes_option_check(self) -> None: + check_selection_options(server_name="Newton", gpu_index=0, force=False) + check_selection_options(server_name="Newton", gpu_index=0, force=True) + check_selection_options(server_name=None, gpu_index=None, force=False) diff --git a/tests/remote/test_observe.py b/tests/remote/test_observe.py new file mode 100644 index 0000000..6a7f8d3 --- /dev/null +++ b/tests/remote/test_observe.py @@ -0,0 +1,113 @@ +"""観測結果の解析と使い方の型のテスト（SPEC_03「取得」「使い方の型」「表示」）.""" + +from __future__ import annotations + +import pytest +from probe_helpers import PMON_HEADER + +from tias_judge_model.remote.observe import ( + GpuSample, + UsagePattern, + classify_usage_pattern, + format_elapsed, + parse_gpu_samples, + parse_pmon, + summarize_samples, +) + + +def _samples(utils: list[int], memory: int = 10) -> list[GpuSample]: + return [GpuSample(util, memory) for util in utils] + + +def _pattern(utils: list[int], process_count: int | None) -> UsagePattern | None: + usage = summarize_samples(_samples(utils)) + assert usage is not None + return classify_usage_pattern(usage, process_count) + + +class TestParseGpuSamples: + def test_groups_samples_by_gpu_index_in_order(self) -> None: + text = "0, 5, 300\n1, 0, 10\n0, 74, 310\n1, 0, 12\n" + samples = parse_gpu_samples(text) + assert samples[0] == [GpuSample(5, 300), GpuSample(74, 310)] + assert samples[1] == [GpuSample(0, 10), GpuSample(0, 12)] + + def test_empty_output_has_no_samples(self) -> None: + assert parse_gpu_samples("") == {} + + +class TestSummarizeSamples: + def test_mean_and_max(self) -> None: + usage = summarize_samples( + [GpuSample(0, 100), GpuSample(74, 300), GpuSample(10, 200), GpuSample(0, 50)] + ) + assert usage is not None + assert usage.sample_count == 4 + assert usage.active_count == 2 + assert usage.mean_utilization_percent == pytest.approx(21.0) + assert usage.max_utilization_percent == 74 + assert usage.max_memory_used_mib == 300 + + def test_no_samples(self) -> None: + assert summarize_samples([]) is None + + +class TestUsagePattern: + @pytest.mark.parametrize( + ("utils", "expected"), + [ + ([1] * 10, UsagePattern.CONSTANT), + ([1] * 8 + [0] * 2, UsagePattern.CONSTANT), # ちょうど 80% + ([1] * 4 + [0], UsagePattern.CONSTANT), # 4/5 = 80% + ([1] * 7 + [0] * 3, UsagePattern.INTERMITTENT), # 70% + ([1] * 79 + [0] * 21, UsagePattern.INTERMITTENT), # 79% + ([1] + [0] * 9, UsagePattern.INTERMITTENT), # 1 件 + ], + ) + def test_active_ratio_boundaries(self, utils: list[int], expected: UsagePattern) -> None: + # 使用率が 0% を超えたサンプルがあれば，計算プロセスの有無によらず常時／断続的 + assert _pattern(utils, 0) is expected + assert _pattern(utils, 2) is expected + + def test_memory_only(self) -> None: + assert _pattern([0, 0, 0], 1) is UsagePattern.MEMORY_ONLY + + def test_idle(self) -> None: + assert _pattern([0, 0, 0], 0) is UsagePattern.IDLE + + +class TestParsePmon: + def test_max_sm_per_compute_process(self) -> None: + text = "\n".join( + [ + *PMON_HEADER, + " 0 1234 C 10 1 - - - - python", + " 1 5678 C 0 0 - - - - python", + " 0 1234 C 74 9 - - - - python", + " 0 1234 C 3 9 - - - - python", + ] + ) + assert parse_pmon(text) == {1234: 74, 5678: 0} + + def test_header_dash_and_graphics_rows_are_excluded(self) -> None: + text = "\n".join( + [ + *PMON_HEADER, + " 0 - - - - - - - - -", + " 0 999 G 50 2 - - - - Xorg", + " 0 1234 C - - - - - - python", + " 0 1234 C 20 - - - - - python", + ] + ) + peaks = parse_pmon(text) + assert peaks == {1234: 20} + + def test_without_header_is_unavailable(self) -> None: + assert parse_pmon("") is None + + +class TestFormatElapsed: + def test_spec_example_hours(self) -> None: + # SPEC_03「表示」の例: 27.7 時間 + assert format_elapsed(99720) == "27.7 時間" diff --git a/tests/remote/test_report.py b/tests/remote/test_report.py new file mode 100644 index 0000000..b707ebc --- /dev/null +++ b/tests/remote/test_report.py @@ -0,0 +1,71 @@ +"""GPU 一覧の表示のテスト（SPEC_03「表示」）. + +各 GPU の状態・使い方の型と使用率の平均／最大・計算プロセスごとの所有ユーザー（自分は「自分」）・ +経過時間・SM 使用率の最大が表示されること． +""" + +from __future__ import annotations + +from probe_helpers import PMON_HEADER, gpu_line, probe_text, sample_lines + +from tias_judge_model.remote.config import GpuFreeThreshold, ServerConfig +from tias_judge_model.remote.gpu import ServerProbe, parse_probe_output, select_gpu +from tias_judge_model.remote.report import format_server_report, format_shared_warning + +THRESHOLD = GpuFreeThreshold( + max_memory_used_mib=1024, max_utilization_percent=80, observe_sec=10, sample_interval_ms=500 +) +UUID0 = "GPU-0000" +SERVER = ServerConfig( + name="Newton", + address="192.168.23.14", + user="hasegawa", + priority=1, + remote_dir="~/work/TIASJudgeModel", + data_dir="~/data/TIASJudgeModel", + data_ready=True, +) + + +def _probe() -> ServerProbe: + text = probe_text( + [gpu_line(0, UUID0, used=500, util=0)], + # 平均 21%・最大 74%（断続的に使用） + samples=sample_lines(0, [0, 74, 10, 0], memory=500), + pmon=[ + *PMON_HEADER, + " 0 100 C 74 1 - - - - python", + " 0 200 C 5 1 - - - - python", + ], + apps=[f"{UUID0}, 100, 250", f"{UUID0}, 200, 250"], + owners=["100 hasegawa 99720", "200 ando 45"], + ) + return ServerProbe(SERVER, parse_probe_output(text, THRESHOLD)) + + +def test_report_shows_state_pattern_users_elapsed_and_sm() -> None: + report = format_server_report([_probe()], show_jobs=False) + assert "Newton" in report + assert "空き（利用者あり）" in report + assert "断続的に使用" in report + assert "平均 21%" in report + assert "最大 74%" in report + assert "自分" in report + assert "ando" in report + assert "27.7 時間" in report + assert "SM 最大 74%" in report + + +def test_report_marks_unreachable_server() -> None: + probe = ServerProbe(SERVER, None, "Connection timed out") + report = format_server_report([probe], show_jobs=True) + assert "Newton" in report + assert "接続不可" in report + + +def test_shared_selection_warns_with_user_names() -> None: + selection = select_gpu([_probe()]) + assert selection is not None + warning = format_shared_warning(selection) + assert warning is not None + assert "ando" in warning diff --git a/tests/remote/test_transfer.py b/tests/remote/test_transfer.py new file mode 100644 index 0000000..3e45428 --- /dev/null +++ b/tests/remote/test_transfer.py @@ -0,0 +1,157 @@ +"""回収した tar の検査・展開のテスト（SPEC_03「pull」の二重の防御の 2 段目）. + +受け取った tar にホワイトリスト外の名前・パス外に出る名前・リンクが 1 つでもあれば， +何も展開せずに中止すること（展開先に何も作られないこと）． +""" + +import io +import tarfile +from pathlib import Path + +import pytest + +from tias_judge_model.remote.files import PushEntry +from tias_judge_model.remote.transfer import extract_pull_archive, write_tar_stream + +RUN_ID = "20260927-120000_newton_abc1234" + + +def _add_file(tar: tarfile.TarFile, name: str, data: bytes = b"x") -> None: + info = tarfile.TarInfo(name) + info.size = len(data) + tar.addfile(info, io.BytesIO(data)) + + +def _add_link(tar: tarfile.TarFile, name: str, target: str, kind: bytes) -> None: + info = tarfile.TarInfo(name) + info.type = kind + info.linkname = target + tar.addfile(info) + + +def _archive(build) -> io.BytesIO: + buffer = io.BytesIO() + with tarfile.open(fileobj=buffer, mode="w", format=tarfile.PAX_FORMAT) as tar: + build(tar) + buffer.seek(0) + return buffer + + +def _all_files(root: Path) -> list[Path]: + return sorted(root.rglob("*")) if root.exists() else [] + + +def test_extracts_whitelisted_files(tmp_path: Path) -> None: + def build(tar: tarfile.TarFile) -> None: + _add_file(tar, f"{RUN_ID}/config.yaml", b"seed: 0\n") + _add_file(tar, f"{RUN_ID}/metrics.json", b"{}") + _add_file(tar, f"{RUN_ID}/figures/confusion_matrix.png", b"png") + + dest = tmp_path / "runs" + result = extract_pull_archive(_archive(build), dest, RUN_ID, with_checkpoints=False) + assert result.files == 3 + assert (dest / RUN_ID / "config.yaml").read_bytes() == b"seed: 0\n" + assert (dest / RUN_ID / "figures" / "confusion_matrix.png").read_bytes() == b"png" + + +def test_overwrites_existing_file(tmp_path: Path) -> None: + dest = tmp_path / "runs" + (dest / RUN_ID).mkdir(parents=True) + (dest / RUN_ID / "metrics.json").write_bytes(b"old") + archive = _archive(lambda tar: _add_file(tar, f"{RUN_ID}/metrics.json", b"new")) + extract_pull_archive(archive, dest, RUN_ID, with_checkpoints=False) + assert (dest / RUN_ID / "metrics.json").read_bytes() == b"new" + + +def test_checkpoints_extracted_only_with_flag(tmp_path: Path) -> None: + def build(tar: tarfile.TarFile) -> None: + _add_file(tar, f"{RUN_ID}/config.yaml") + _add_file(tar, f"{RUN_ID}/checkpoints/best.pt", b"weights") + + dest = tmp_path / "runs" + with pytest.raises(ValueError): + extract_pull_archive(_archive(build), dest, RUN_ID, with_checkpoints=False) + assert _all_files(dest) == [] + + extract_pull_archive(_archive(build), dest, RUN_ID, with_checkpoints=True) + assert (dest / RUN_ID / "checkpoints" / "best.pt").read_bytes() == b"weights" + + +BAD_MEMBERS = { + "split_json": lambda tar: _add_file(tar, f"{RUN_ID}/split.json"), + "predictions": lambda tar: _add_file(tar, f"{RUN_ID}/predictions.csv"), + "png_outside_figures": lambda tar: _add_file(tar, f"{RUN_ID}/P0001.png"), + "manifest": lambda tar: _add_file(tar, f"{RUN_ID}/manifest.csv"), + "image": lambda tar: _add_file(tar, f"{RUN_ID}/images/P0001.jpg"), + "other_run": lambda tar: _add_file(tar, "other_run/config.yaml"), + "dotdot": lambda tar: _add_file(tar, f"{RUN_ID}/../evil.txt"), + "dotdot_top": lambda tar: _add_file(tar, "../evil.txt"), + "absolute": lambda tar: _add_file(tar, "/tmp/evil.txt"), + "drive_letter": lambda tar: _add_file(tar, "C:/evil.txt"), + "drive_relative": lambda tar: _add_file(tar, "C:evil.txt"), + "backslash": lambda tar: _add_file(tar, f"{RUN_ID}\\config.yaml"), + "symlink": lambda tar: _add_link(tar, f"{RUN_ID}/log.txt", "/etc/passwd", tarfile.SYMTYPE), + "symlink_relative": lambda tar: _add_link( + tar, f"{RUN_ID}/history.csv", "../../outside", tarfile.SYMTYPE + ), + "hardlink": lambda tar: _add_link( + tar, f"{RUN_ID}/metrics.json", f"{RUN_ID}/config.yaml", tarfile.LNKTYPE + ), +} + + +@pytest.mark.parametrize("bad", list(BAD_MEMBERS), ids=list(BAD_MEMBERS)) +@pytest.mark.parametrize("position", ["first", "last"]) +def test_bad_member_aborts_without_extracting_anything( + tmp_path: Path, bad: str, position: str +) -> None: + def build(tar: tarfile.TarFile) -> None: + if position == "first": + BAD_MEMBERS[bad](tar) + _add_file(tar, f"{RUN_ID}/config.yaml") + _add_file(tar, f"{RUN_ID}/figures/cm.png") + if position == "last": + BAD_MEMBERS[bad](tar) + + work = tmp_path / "work" + dest = work / "runs" + with pytest.raises(ValueError): + extract_pull_archive(_archive(build), dest, RUN_ID, with_checkpoints=True) + # 展開先にも，その外側（パス外に出る名前の行き先）にも何も作られない + assert _all_files(dest) == [] + assert not (dest / RUN_ID / "config.yaml").exists() + assert not (work / "evil.txt").exists() + assert not (tmp_path / "evil.txt").exists() + assert [p for p in _all_files(tmp_path) if p.is_file()] == [] + + +def test_error_message_does_not_contain_rejected_names(tmp_path: Path) -> None: + secret = "P0001_yamada" + + def build(tar: tarfile.TarFile) -> None: + _add_file(tar, f"{RUN_ID}/config.yaml") + _add_file(tar, f"{RUN_ID}/{secret}.png") + + with pytest.raises(ValueError) as excinfo: + extract_pull_archive(_archive(build), tmp_path / "runs", RUN_ID, with_checkpoints=False) + assert secret not in str(excinfo.value) + + +def test_write_tar_stream_round_trip(tmp_path: Path) -> None: + src = tmp_path / "src" + (src / "pkg").mkdir(parents=True) + (src / "pkg" / "a.py").write_bytes(b"print('a')\n") + (src / "日本語.md").write_text("x", encoding="utf-8") + entries = [ + PushEntry(src / "pkg" / "a.py", "pkg/a.py", (src / "pkg" / "a.py").stat().st_size), + PushEntry(src / "日本語.md", "日本語.md", 1), + ] + buffer = io.BytesIO() + write_tar_stream(entries, buffer) + buffer.seek(0) + with tarfile.open(fileobj=buffer, mode="r") as tar: + names = tar.getnames() + assert names == ["pkg/a.py", "日本語.md"] + member = tar.extractfile("pkg/a.py") + assert member is not None + assert member.read() == b"print('a')\n" diff --git a/tests/remote/test_whitelist.py b/tests/remote/test_whitelist.py new file mode 100644 index 0000000..9854ce6 --- /dev/null +++ b/tests/remote/test_whitelist.py @@ -0,0 +1,194 @@ +"""回収ホワイトリストの判定のテスト（SPEC_01「回収ホワイトリスト」・SPEC_03「pull」）. + +最重要: 名前が一致するファイルだけが回収対象になり，分割結果・個票の予測・画像・manifest 等は +回収されないこと．testing ルール必須対象 4． +""" + +import io +import tarfile + +import pytest + +from tias_judge_model.remote.whitelist import ( + check_pull_members, + is_pull_allowed, + parse_find_listing, + plan_pull, + validate_run_id, +) + +RUN_ID = "20260927-120000_newton_abc1234" + +ROOT_FILES = ["config.yaml", "history.csv", "metrics.json", "confusion_matrix.csv", "log.txt"] + +# 列挙外（回収してはならない）ファイル．患者 ID・画像名を含みうる名前も混ぜる +NOT_ALLOWED = [ + "split.json", + "predictions.csv", + "per_image_predictions.csv", + "manifest.csv", + "confusion_matrix.png", # figures/ 以外にある *.png + "P0001_tongue.png", + "P0001_tongue.jpg", + "images/pale_1/P0001.png", + "misclassified/P0001.png", + "data/manifest.csv", + "config.yml", + "config.yaml.bak", + "metrics.json.tmp", + "sub/config.yaml", + "figures/config.yaml", + "figures/cm.jpg", + "figures/sub/cm.png", + "checkpoints/best.pth", + "checkpoints/sub/best.pt", + "best.pt", +] + + +class TestIsPullAllowed: + @pytest.mark.parametrize("name", ROOT_FILES) + def test_root_files_are_allowed(self, name: str) -> None: + assert is_pull_allowed(name, with_checkpoints=False) + + @pytest.mark.parametrize("name", ["figures/confusion_matrix.png", "figures/learning_curve.png"]) + def test_figures_png_is_allowed(self, name: str) -> None: + assert is_pull_allowed(name, with_checkpoints=False) + + @pytest.mark.parametrize("name", NOT_ALLOWED) + def test_names_outside_whitelist_are_not_allowed(self, name: str) -> None: + assert not is_pull_allowed(name, with_checkpoints=False) + # --with-checkpoints は checkpoints/*.pt 以外の判定を変えない + assert not is_pull_allowed(name, with_checkpoints=True) + + def test_checkpoints_only_with_flag(self) -> None: + assert not is_pull_allowed("checkpoints/best.pt", with_checkpoints=False) + assert is_pull_allowed("checkpoints/best.pt", with_checkpoints=True) + + @pytest.mark.parametrize( + "name", + [ + "../config.yaml", + "figures/../config.yaml", + "/config.yaml", + "C:config.yaml", + "C:/config.yaml", + "figures\\cm.png", + "..\\config.yaml", + ], + ) + def test_unsafe_paths_are_not_allowed(self, name: str) -> None: + assert not is_pull_allowed(name, with_checkpoints=True) + + +class TestValidateRunId: + def test_accepts_normal_run_id(self) -> None: + assert validate_run_id(RUN_ID) == RUN_ID + + @pytest.mark.parametrize( + "run_id", ["a/b", "a\\b", "..", "../x", "x/..", "a..b", "/abs", "runs/x", ""] + ) + def test_rejects_run_id_with_separator_or_dotdot(self, run_id: str) -> None: + with pytest.raises(ValueError): + validate_run_id(run_id) + + +class TestPlanPull: + def _listing(self, names: list[str], kind: str = "f") -> list[tuple[str, str]]: + return [(kind, f"{RUN_ID}/{name}") for name in names] + + def test_selects_only_whitelisted_files(self) -> None: + allowed = [*ROOT_FILES, "figures/cm.png"] + listing = self._listing(allowed + NOT_ALLOWED + ["checkpoints/best.pt"]) + plan = plan_pull(listing, RUN_ID, with_checkpoints=False) + assert sorted(plan.files) == sorted(f"{RUN_ID}/{name}" for name in allowed) + assert plan.skipped == len(NOT_ALLOWED) + 1 + + def test_with_checkpoints_includes_pt(self) -> None: + listing = self._listing(["config.yaml", "checkpoints/best.pt", "split.json"]) + plan = plan_pull(listing, RUN_ID, with_checkpoints=True) + assert sorted(plan.files) == [f"{RUN_ID}/checkpoints/best.pt", f"{RUN_ID}/config.yaml"] + assert plan.skipped == 1 + + def test_links_are_not_selected_even_with_whitelisted_name(self) -> None: + listing = self._listing(["config.yaml"], kind="l") + plan = plan_pull(listing, RUN_ID, with_checkpoints=False) + assert plan.files == () + assert plan.skipped == 1 + + def test_files_of_other_run_are_not_selected(self) -> None: + plan = plan_pull([("f", "other_run/config.yaml")], RUN_ID, with_checkpoints=False) + assert plan.files == () + + def test_parse_find_listing_round_trip(self) -> None: + data = f"f {RUN_ID}/config.yaml\0f {RUN_ID}/figures/a b.png\0l {RUN_ID}/x\0".encode() + assert parse_find_listing(data) == [ + ("f", f"{RUN_ID}/config.yaml"), + ("f", f"{RUN_ID}/figures/a b.png"), + ("l", f"{RUN_ID}/x"), + ] + + +def _regular(name: str) -> tarfile.TarInfo: + info = tarfile.TarInfo(name) + info.size = 0 + return info + + +def _link(name: str, target: str, kind: bytes) -> tarfile.TarInfo: + info = tarfile.TarInfo(name) + info.type = kind + info.linkname = target + return info + + +class TestCheckPullMembers: + def test_accepts_only_whitelisted_regular_files(self) -> None: + members = [_regular(f"{RUN_ID}/{name}") for name in [*ROOT_FILES, "figures/cm.png"]] + assert check_pull_members(members, RUN_ID, with_checkpoints=False) == len(members) + + @pytest.mark.parametrize( + "bad", + [ + *(_regular(f"{RUN_ID}/{name}") for name in NOT_ALLOWED), + _regular(f"{RUN_ID}/checkpoints/best.pt"), + _regular("other_run/config.yaml"), + _regular("config.yaml"), + _regular(f"{RUN_ID}/../config.yaml"), + _regular(f"/{RUN_ID}/config.yaml"), + _regular(f"C:/{RUN_ID}/config.yaml"), + _regular(f"{RUN_ID}\\config.yaml"), + _link(f"{RUN_ID}/config.yaml", "/etc/passwd", tarfile.SYMTYPE), + _link(f"{RUN_ID}/metrics.json", f"{RUN_ID}/config.yaml", tarfile.LNKTYPE), + ], + ) + def test_rejects_whole_archive_if_any_member_is_bad(self, bad: tarfile.TarInfo) -> None: + members = [_regular(f"{RUN_ID}/config.yaml"), bad, _regular(f"{RUN_ID}/log.txt")] + with pytest.raises(ValueError): + check_pull_members(members, RUN_ID, with_checkpoints=False) + + def test_error_message_does_not_contain_member_names(self) -> None: + secret = "P0001_yamada_tongue" + members = [ + _regular(f"{RUN_ID}/{secret}.png"), + _link(f"{RUN_ID}/{secret}_link", "/etc/passwd", tarfile.SYMTYPE), + _regular(f"../{secret}.csv"), + ] + with pytest.raises(ValueError) as excinfo: + check_pull_members(members, RUN_ID, with_checkpoints=False) + assert secret not in str(excinfo.value) + + def test_checkpoints_accepted_only_with_flag(self) -> None: + members = [_regular(f"{RUN_ID}/checkpoints/best.pt")] + assert check_pull_members(members, RUN_ID, with_checkpoints=True) == 1 + with pytest.raises(ValueError): + check_pull_members(members, RUN_ID, with_checkpoints=False) + + def test_members_from_real_tar_file(self) -> None: + buffer = io.BytesIO() + with tarfile.open(fileobj=buffer, mode="w") as tar: + tar.addfile(_regular(f"{RUN_ID}/config.yaml"), io.BytesIO(b"")) + tar.addfile(_link(f"{RUN_ID}/log.txt", "/etc/shadow", tarfile.SYMTYPE)) + buffer.seek(0) + with tarfile.open(fileobj=buffer, mode="r") as tar, pytest.raises(ValueError): + check_pull_members(tar.getmembers(), RUN_ID, with_checkpoints=False) diff --git a/uv.lock b/uv.lock index c1a2756..0c8df43 100644 --- a/uv.lock +++ b/uv.lock @@ -1,6 +1,10 @@ version = 1 revision = 3 requires-python = "==3.12.*" +resolution-markers = [ + "sys_platform == 'linux'", + "sys_platform != 'linux'", +] [[package]] name = "annotated-types" @@ -75,14 +79,14 @@ [[package]] name = "cuda-bindings" -version = "13.4.3" +version = "12.9.9" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cuda-pathfinder" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/65/11/1293429c1c3a3e19b551275e65efddd122a905bbe7e368816a59f3ef2a41/cuda_bindings-13.4.3-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bfbd3f7d4ac04dd41dc49121b9e408c8283992f47124c2290ecb79bbbadcca8e", size = 6488226, upload-time = "2026-09-23T02:22:00.578Z" }, - { url = "https://files.pythonhosted.org/packages/b8/c3/efb6bbb7307bf5c83dc4acca650280b210c67ed1a1a60f898a90e7c82e38/cuda_bindings-13.4.3-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d5f72bcfcdf3be23e1da3c792f68f508586f48d037bca8b10f552c4cca5971f2", size = 7172185, upload-time = "2026-09-23T02:22:03.114Z" }, + { url = "https://files.pythonhosted.org/packages/15/d5/572ced1aecc8c8f3e617c3af58d3b1ed5378ed71531c607c960a3e48b516/cuda_bindings-12.9.9-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:88c41afe5be41acb17e57977dd534e335f86d0c96d284d71457168525f6327c0", size = 7063842, upload-time = "2026-09-23T03:20:45.94Z" }, + { url = "https://files.pythonhosted.org/packages/96/e3/62c521e36276008c5e53798e7751ca33e16a16d05f517426c2ffb4356e47/cuda_bindings-12.9.9-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fb842851be2ed9e80fce988dd68c8ea27a506a71e0e29d228b35cbd0a333dac2", size = 7627343, upload-time = "2026-09-23T03:20:48.555Z" }, ] [[package]] @@ -95,51 +99,42 @@ [[package]] name = "cuda-toolkit" -version = "13.0.3.0" +version = "12.6.3" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d1/c7/a79086a62c98befcdb8349656c6f114e2db3b8b2422f6e25c97a7f2a9a3c/cuda_toolkit-13.0.3.0-py2.py3-none-any.whl", hash = "sha256:d693caaa261214ddd7dbb60d68e71cbed884e68c2be7509778f3051da0b91c3f", size = 2512, upload-time = "2026-04-14T00:50:08.173Z" }, + { url = "https://files.pythonhosted.org/packages/ad/88/2dbc37975fffb874418b14380418a1b99cb36f2101fd1d08c54e06ee8c95/cuda_toolkit-12.6.3-py2.py3-none-any.whl", hash = "sha256:79d8605baeb6c2f695761e0efb54bc62dbc3c9e32eb0742df7669c07befaa8f7", size = 2288, upload-time = "2025-08-13T02:03:05.283Z" }, ] [package.optional-dependencies] cublas = [ - { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, - { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cublas-cu12" }, ] cudart = [ - { name = "nvidia-cuda-runtime", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cuda-runtime-cu12" }, ] cufft = [ - { name = "nvidia-cufft", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, - { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cufft-cu12" }, ] cufile = [ - { name = "nvidia-cufile", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cufile-cu12" }, ] cupti = [ - { name = "nvidia-cuda-cupti", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cuda-cupti-cu12" }, ] curand = [ - { name = "nvidia-curand", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-curand-cu12" }, ] cusolver = [ - { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, - { name = "nvidia-cusolver", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, - { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, - { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cusolver-cu12" }, ] cusparse = [ - { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, - { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, -] -nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cusparse-cu12" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cuda-nvrtc-cu12" }, ] nvtx = [ - { name = "nvidia-nvtx", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvtx-cu12" }, ] [[package]] @@ -430,155 +425,166 @@ ] [[package]] -name = "nvidia-cublas" -version = "13.1.1.3" +name = "nvidia-cublas-cu12" +version = "12.6.4.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/af/eb/ff4b8c503fa1f1796679dce648854d58751982426e4e4b37d6fce49d259c/nvidia_cublas_cu12-12.6.4.1-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08ed2686e9875d01b58e3cb379c6896df8e76c75e0d4a7f7dace3d7b6d9ef8eb", size = 393138322, upload-time = "2024-11-20T17:40:25.65Z" }, + { url = "https://files.pythonhosted.org/packages/97/0d/f1f0cadbf69d5b9ef2e4f744c9466cb0a850741d08350736dfdb4aa89569/nvidia_cublas_cu12-12.6.4.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:235f728d6e2a409eddf1df58d5b0921cf80cfa9e72b9f2775ccb7b4a87984668", size = 390794615, upload-time = "2024-11-20T17:39:52.715Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.6.80" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e6/8b/2f6230cb715646c3a9425636e513227ce5c93c4d65823a734f4bb86d43c3/nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:166ee35a3ff1587f2490364f90eeeb8da06cd867bd5b701bf7f9a02b78bc63fc", size = 8236764, upload-time = "2024-11-20T17:35:41.03Z" }, + { url = "https://files.pythonhosted.org/packages/25/0f/acb326ac8fd26e13c799e0b4f3b2751543e1834f04d62e729485872198d4/nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_aarch64.whl", hash = "sha256:358b4a1d35370353d52e12f0a7d1769fc01ff74a191689d3870b2123156184c4", size = 8236756, upload-time = "2024-10-01T16:57:45.507Z" }, + { url = "https://files.pythonhosted.org/packages/49/60/7b6497946d74bcf1de852a21824d63baad12cd417db4195fc1bfe59db953/nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6768bad6cab4f19e8292125e5f1ac8aa7d1718704012a0e3272a6f61c4bce132", size = 8917980, upload-time = "2024-11-20T17:36:04.019Z" }, + { url = "https://files.pythonhosted.org/packages/a5/24/120ee57b218d9952c379d1e026c4479c9ece9997a4fb46303611ee48f038/nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a3eff6cdfcc6a4c35db968a06fcadb061cbc7d6dde548609a941ff8701b98b73", size = 8917972, upload-time = "2024-10-01T16:58:06.036Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.6.85" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f5/31/ffb400c5ae99daf09687aa6c42831c5d824f71c4851363ed2a4a1ac52bab/nvidia_cuda_nvrtc_cu12-12.6.85-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:800927308ccc5dd6246d3f61f7fcef2ed7ec4e59e199090d360d3293f78bd5a2", size = 23649944, upload-time = "2024-11-20T17:38:06.511Z" }, + { url = "https://files.pythonhosted.org/packages/48/ab/476146f59ff5ef5bd6e62c187097d859ea78b5752d19c6c3f9be5f90dafc/nvidia_cuda_nvrtc_cu12-12.6.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3f3134f50963882373063901657554f230bedf6039d30b09f6be55c64c993a37", size = 23162872, upload-time = "2024-11-20T17:37:42.967Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.6.77" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8f/ea/590b2ac00d772a8abd1c387a92b46486d2679ca6622fd25c18ff76265663/nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6116fad3e049e04791c0256a9778c16237837c08b27ed8c8401e2e45de8d60cd", size = 908052, upload-time = "2024-11-20T17:35:19.905Z" }, + { url = "https://files.pythonhosted.org/packages/b7/3d/159023799677126e20c8fd580cca09eeb28d5c5a624adc7f793b9aa8bbfa/nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:d461264ecb429c84c8879a7153499ddc7b19b5f8d84c204307491989a365588e", size = 908040, upload-time = "2024-10-01T16:57:22.221Z" }, + { url = "https://files.pythonhosted.org/packages/e1/23/e717c5ac26d26cf39a27fbc076240fad2e3b817e5889d671b67f4f9f49c5/nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ba3b56a4f896141e25e19ab287cd71e52a6a0f4b29d0d31609f60e3b4d5219b7", size = 897690, upload-time = "2024-11-20T17:35:30.697Z" }, + { url = "https://files.pythonhosted.org/packages/f0/62/65c05e161eeddbafeca24dc461f47de550d9fa8a7e04eb213e32b55cfd99/nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a84d15d5e1da416dd4774cb42edf5e954a3e60cc945698dc1d5be02321c44dc8", size = 897678, upload-time = "2024-10-01T16:57:33.821Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.10.2.21" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cuda-nvrtc" }, + { name = "nvidia-cublas-cu12" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" }, - { url = "https://files.pythonhosted.org/packages/3b/cd/154ca20c38269e05eff77c1464e6c1da89f50a6390b565e9d82e06bc11e1/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:37936a16db8fe4ac1f065c2139360608a543a09275cb1a1af612e08cfa065436", size = 423138758, upload-time = "2026-04-08T18:46:58.655Z" }, + { url = "https://files.pythonhosted.org/packages/fa/41/e79269ce215c857c935fd86bcfe91a451a584dfc27f1e068f568b9ad1ab7/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:c9132cc3f8958447b4910a1720036d9eff5928cc3179b0a51fb6d167c6cc87d8", size = 705026878, upload-time = "2025-06-06T21:52:51.348Z" }, + { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" }, ] [[package]] -name = "nvidia-cuda-cupti" -version = "13.0.85" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2a/2a/80353b103fc20ce05ef51e928daed4b6015db4aaa9162ed0997090fe2250/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151", size = 10310827, upload-time = "2025-09-04T08:26:42.012Z" }, - { url = "https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8", size = 10715597, upload-time = "2025-09-04T08:26:51.312Z" }, -] - -[[package]] -name = "nvidia-cuda-nvrtc" -version = "13.0.88" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575", size = 90215200, upload-time = "2025-09-04T08:28:44.204Z" }, - { url = "https://files.pythonhosted.org/packages/b7/dc/6bb80850e0b7edd6588d560758f17e0550893a1feaf436807d64d2da040f/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b", size = 43015449, upload-time = "2025-09-04T08:28:20.239Z" }, -] - -[[package]] -name = "nvidia-cuda-runtime" -version = "13.0.96" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/87/4f/17d7b9b8e285199c58ce28e31b5c5bbaa4d8271af06a89b6405258245de2/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55", size = 2261060, upload-time = "2025-10-09T08:55:15.78Z" }, - { url = "https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548", size = 2243632, upload-time = "2025-10-09T08:55:36.117Z" }, -] - -[[package]] -name = "nvidia-cudnn-cu13" -version = "9.24.0.43" +name = "nvidia-cufft-cu12" +version = "11.3.0.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas" }, + { name = "nvidia-nvjitlink-cu12" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/ca/30/7c257e3d5cb4fecb147b93895c66e29c93f8e76d74b45bb418ff0587c4ec/nvidia_cudnn_cu13-9.24.0.43-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:a6812a554a1ff0413e9c52b84c26c050380649ab9615f9c16bded368ce9f421f", size = 650976863, upload-time = "2026-07-02T16:23:39.248Z" }, - { url = "https://files.pythonhosted.org/packages/5c/ba/791cffd048fe5b044e620df55267e3e95c0e6e07d50b41e377c03dfc910f/nvidia_cudnn_cu13-9.24.0.43-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:71f181cd810e90f9b6023b01186fe82d13d65f0ec098581ee201d39fad769e4b", size = 553099438, upload-time = "2026-07-02T16:27:42.58Z" }, + { url = "https://files.pythonhosted.org/packages/1f/37/c50d2b2f2c07e146776389e3080f4faf70bcc4fa6e19d65bb54ca174ebc3/nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d16079550df460376455cba121db6564089176d9bac9e4f360493ca4741b22a6", size = 200164144, upload-time = "2024-11-20T17:40:58.288Z" }, + { url = "https://files.pythonhosted.org/packages/ce/f5/188566814b7339e893f8d210d3a5332352b1409815908dad6a363dcceac1/nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_aarch64.whl", hash = "sha256:8510990de9f96c803a051822618d42bf6cb8f069ff3f48d93a8486efdacb48fb", size = 200164135, upload-time = "2024-10-01T17:03:24.212Z" }, + { url = "https://files.pythonhosted.org/packages/8f/16/73727675941ab8e6ffd86ca3a4b7b47065edcca7a997920b831f8147c99d/nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ccba62eb9cef5559abd5e0d54ceed2d9934030f51163df018532142a8ec533e5", size = 200221632, upload-time = "2024-11-20T17:41:32.357Z" }, + { url = "https://files.pythonhosted.org/packages/60/de/99ec247a07ea40c969d904fc14f3a356b3e2a704121675b75c366b694ee1/nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_x86_64.whl", hash = "sha256:768160ac89f6f7b459bee747e8d175dbf53619cfe74b2a5636264163138013ca", size = 200221622, upload-time = "2024-10-01T17:03:58.79Z" }, ] [[package]] -name = "nvidia-cufft" -version = "12.0.0.61" +name = "nvidia-cufile-cu12" +version = "1.11.1.6" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b2/66/cc9876340ac68ae71b15c743ddb13f8b30d5244af344ec8322b449e35426/nvidia_cufile_cu12-1.11.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:cc23469d1c7e52ce6c1d55253273d32c565dd22068647f3aa59b3c6b005bf159", size = 1142103, upload-time = "2024-11-20T17:42:11.83Z" }, + { url = "https://files.pythonhosted.org/packages/17/bf/cc834147263b929229ce4aadd62869f0b195e98569d4c28b23edc72b85d9/nvidia_cufile_cu12-1.11.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:8f57a0051dcf2543f6dc2b98a98cb2719c37d3cee1baba8965d57f3bbc90d4db", size = 1066155, upload-time = "2024-11-20T17:41:49.376Z" }, +] + +[[package]] +name = "nvidia-curand-cu12" +version = "10.3.7.77" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/42/ac/36543605358a355632f1a6faa3e2d5dfb91eab1e4bc7d552040e0383c335/nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:6e82df077060ea28e37f48a3ec442a8f47690c7499bff392a5938614b56c98d8", size = 56289881, upload-time = "2024-10-01T17:04:18.981Z" }, + { url = "https://files.pythonhosted.org/packages/73/1b/44a01c4e70933637c93e6e1a8063d1e998b50213a6b65ac5a9169c47e98e/nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a42cd1344297f70b9e39a1e4f467a4e1c10f1da54ff7a85c12197f6c652c8bdf", size = 56279010, upload-time = "2024-11-20T17:42:50.958Z" }, + { url = "https://files.pythonhosted.org/packages/4a/aa/2c7ff0b5ee02eaef890c0ce7d4f74bc30901871c5e45dee1ae6d0083cd80/nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:99f1a32f1ac2bd134897fc7a203f779303261268a65762a623bf30cc9fe79117", size = 56279000, upload-time = "2024-10-01T17:04:45.274Z" }, + { url = "https://files.pythonhosted.org/packages/a6/02/5362a9396f23f7de1dd8a64369e87c85ffff8216fc8194ace0fa45ba27a5/nvidia_curand_cu12-10.3.7.77-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:7b2ed8e95595c3591d984ea3603dd66fe6ce6812b886d59049988a712ed06b6e", size = 56289882, upload-time = "2024-11-20T17:42:25.222Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.7.1.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink" }, + { name = "nvidia-cublas-cu12" }, + { name = "nvidia-cusparse-cu12" }, + { name = "nvidia-nvjitlink-cu12" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, - { url = "https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3", size = 214085489, upload-time = "2025-09-04T08:31:56.044Z" }, + { url = "https://files.pythonhosted.org/packages/93/17/dbe1aa865e4fdc7b6d4d0dd308fdd5aaab60f939abfc0ea1954eac4fb113/nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_aarch64.whl", hash = "sha256:0ce237ef60acde1efc457335a2ddadfd7610b892d94efee7b776c64bb1cac9e0", size = 157833628, upload-time = "2024-10-01T17:05:05.591Z" }, + { url = "https://files.pythonhosted.org/packages/f0/6e/c2cf12c9ff8b872e92b4a5740701e51ff17689c4d726fca91875b07f655d/nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e9e49843a7707e42022babb9bcfa33c29857a93b88020c4e4434656a655b698c", size = 158229790, upload-time = "2024-11-20T17:43:43.211Z" }, + { url = "https://files.pythonhosted.org/packages/9f/81/baba53585da791d043c10084cf9553e074548408e04ae884cfe9193bd484/nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:6cf28f17f64107a0c4d7802be5ff5537b2130bfc112f25d5a30df227058ca0e6", size = 158229780, upload-time = "2024-10-01T17:05:39.875Z" }, + { url = "https://files.pythonhosted.org/packages/7c/5f/07d0ba3b7f19be5a5ec32a8679fc9384cfd9fc6c869825e93be9f28d6690/nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:dbbe4fc38ec1289c7e5230e16248365e375c3673c9c8bac5796e2e20db07f56e", size = 157833630, upload-time = "2024-11-20T17:43:16.77Z" }, ] [[package]] -name = "nvidia-cufile" -version = "1.15.1.6" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44", size = 1223672, upload-time = "2025-09-04T08:32:22.779Z" }, - { url = "https://files.pythonhosted.org/packages/ab/73/cc4a14c9813a8a0d509417cf5f4bdaba76e924d58beb9864f5a7baceefbf/nvidia_cufile-1.15.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1", size = 1136992, upload-time = "2025-09-04T08:32:14.119Z" }, -] - -[[package]] -name = "nvidia-curand" -version = "10.4.0.35" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1e/72/7c2ae24fb6b63a32e6ae5d241cc65263ea18d08802aaae087d9f013335a2/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a", size = 61962106, upload-time = "2025-08-04T10:21:41.128Z" }, - { url = "https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc", size = 59544258, upload-time = "2025-08-04T10:22:03.992Z" }, -] - -[[package]] -name = "nvidia-cusolver" -version = "12.0.4.66" +name = "nvidia-cusparse-cu12" +version = "12.5.4.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas" }, - { name = "nvidia-cusparse" }, - { name = "nvidia-nvjitlink" }, + { name = "nvidia-nvjitlink-cu12" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, - { url = "https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112", size = 200941980, upload-time = "2025-09-04T08:33:22.767Z" }, + { url = "https://files.pythonhosted.org/packages/eb/eb/6681efd0aa7df96b4f8067b3ce7246833dd36830bb4cec8896182773db7d/nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d25b62fb18751758fe3c93a4a08eff08effedfe4edf1c6bb5afd0890fe88f887", size = 216451147, upload-time = "2024-11-20T17:44:18.055Z" }, + { url = "https://files.pythonhosted.org/packages/d3/56/3af21e43014eb40134dea004e8d0f1ef19d9596a39e4d497d5a7de01669f/nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_aarch64.whl", hash = "sha256:7aa32fa5470cf754f72d1116c7cbc300b4e638d3ae5304cfa4a638a5b87161b1", size = 216451135, upload-time = "2024-10-01T17:06:03.826Z" }, + { url = "https://files.pythonhosted.org/packages/06/1e/b8b7c2f4099a37b96af5c9bb158632ea9e5d9d27d7391d7eb8fc45236674/nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7556d9eca156e18184b94947ade0fba5bb47d69cec46bf8660fd2c71a4b48b73", size = 216561367, upload-time = "2024-11-20T17:44:54.824Z" }, + { url = "https://files.pythonhosted.org/packages/43/ac/64c4316ba163e8217a99680c7605f779accffc6a4bcd0c778c12948d3707/nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:23749a6571191a215cb74d1cdbff4a86e7b19f1200c071b3fcf844a5bea23a2f", size = 216561357, upload-time = "2024-10-01T17:06:29.861Z" }, ] [[package]] -name = "nvidia-cusparse" -version = "12.6.3.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-nvjitlink" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, - { url = "https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b", size = 145942937, upload-time = "2025-09-04T08:33:58.029Z" }, -] - -[[package]] -name = "nvidia-cusparselt-cu13" -version = "0.8.1" +name = "nvidia-cusparselt-cu12" +version = "0.7.1" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/46/e1/cdc1797eadf82d3a9a575a19b33fdc871a97edbec42c00b5b5e914f4aff4/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f", size = 221051344, upload-time = "2025-09-05T18:49:51.289Z" }, - { url = "https://files.pythonhosted.org/packages/34/7d/2661f2fb3ac4302f3a246f5fc030213ac60c1fe0bce84f9783dbd831dbb7/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:786ce87568c303fadb5afcc7102d454cd3040d75f6f8626f5db460d1871f4dd0", size = 170148586, upload-time = "2025-09-05T18:50:50.248Z" }, + { url = "https://files.pythonhosted.org/packages/73/b9/598f6ff36faaece4b3c50d26f50e38661499ff34346f00e057760b35cc9d/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:8878dce784d0fac90131b6817b607e803c36e629ba34dc5b433471382196b6a5", size = 283835557, upload-time = "2025-02-26T00:16:54.265Z" }, + { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" }, ] [[package]] -name = "nvidia-nccl-cu13" -version = "2.30.7" +name = "nvidia-nccl-cu12" +version = "2.29.3" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d1/21/a73174c6157101bdf1ffc22b517f76ff0082613989dd9bc8f43e8034caac/nvidia_nccl_cu13-2.30.7-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:ca786ffa5a647c75d4d1f5cc72a6c4f537947e2ba8823d7c8aaf768e7a7b9f77", size = 215983881, upload-time = "2026-06-09T03:23:15.633Z" }, - { url = "https://files.pythonhosted.org/packages/3f/34/c500f90c7ae641b8e0f98965b36b8a7ac79cc8b296e8d251fe3eb592ee54/nvidia_nccl_cu13-2.30.7-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:cefa7fdb9710efd0f39c5f1be1d61ff6fc9a996c451265bd7fbdcf9455ed4b50", size = 215965170, upload-time = "2026-06-09T03:23:39.73Z" }, + { url = "https://files.pythonhosted.org/packages/28/cf/bcf8bb0c0030b1b9a345331f6281c37d2a8669758521eb93c382f6f87c8f/nvidia_nccl_cu12-2.29.3-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:6351b79dc7d2cc3d654ea1523616b9eeded71fe9c8da66b71eef9a5d1b2adad4", size = 289708535, upload-time = "2026-02-03T21:10:58.804Z" }, + { url = "https://files.pythonhosted.org/packages/31/5a/cac7d231f322b66caa16fd4b136ebc8e4b18b2805811c2d58dc47210cdea/nvidia_nccl_cu12-2.29.3-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:35ad42e7d5d722a83c36a3a478e281c20a5646383deaf1b9ed1a9ab7d61bed53", size = 289760316, upload-time = "2026-02-03T21:11:37.899Z" }, ] [[package]] -name = "nvidia-nvjitlink" -version = "13.4.92" +name = "nvidia-nvjitlink-cu12" +version = "12.9.86" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1d/6b/eef7a9e32872b8f41e145bf10cddc9af26e153c338852811fe9a9baddf9e/nvidia_nvjitlink-13.4.92-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:e0391f24ed94ec879b84e3da4d4ec320c879aff681f2c7a638462f7199284323", size = 42452378, upload-time = "2026-09-16T20:45:29.042Z" }, - { url = "https://files.pythonhosted.org/packages/1f/a8/1cbd4014898af8b419e69b0d7dbc63da2121ee92d92b47d59f4fe9075349/nvidia_nvjitlink-13.4.92-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:25f74fad0d654271c921ac4dca614bd6258bc21791242fc7b2289dad7ae9c099", size = 40420120, upload-time = "2026-09-16T20:45:19.163Z" }, + { url = "https://files.pythonhosted.org/packages/46/0c/c75bbfb967457a0b7670b8ad267bfc4fffdf341c074e0a80db06c24ccfd4/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:e3f1171dbdc83c5932a45f0f4c99180a70de9bd2718c1ab77d14104f6d7147f9", size = 39748338, upload-time = "2025-06-05T20:10:25.613Z" }, + { url = "https://files.pythonhosted.org/packages/97/bc/2dcba8e70cf3115b400fef54f213bcd6715a3195eba000f8330f11e40c45/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:994a05ef08ef4b0b299829cde613a424382aff7efb08a7172c1fa616cc3af2ca", size = 39514880, upload-time = "2025-06-05T20:10:04.89Z" }, ] [[package]] -name = "nvidia-nvshmem-cu13" +name = "nvidia-nvshmem-cu12" version = "3.4.5" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/dc/0f/05cc9c720236dcd2db9c1ab97fff629e96821be2e63103569da0c9b72f19/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dc2a197f38e5d0376ad52cd1a2a3617d3cdc150fd5966f4aee9bcebb1d68fe9", size = 60215947, upload-time = "2025-09-06T00:32:20.022Z" }, - { url = "https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80", size = 60412546, upload-time = "2025-09-06T00:32:41.564Z" }, + { url = "https://files.pythonhosted.org/packages/1d/6a/03aa43cc9bd3ad91553a88b5f6fb25ed6a3752ae86ce2180221962bc2aa5/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0b48363fc6964dede448029434c6abed6c5e37f823cb43c3bcde7ecfc0457e15", size = 138936938, upload-time = "2025-09-06T00:32:05.589Z" }, + { url = "https://files.pythonhosted.org/packages/b5/09/6ea3ea725f82e1e76684f0708bbedd871fc96da89945adeba65c3835a64c/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:042f2500f24c021db8a06c5eec2539027d57460e1c1a762055a6554f72c369bd", size = 139103095, upload-time = "2025-09-06T00:32:31.266Z" }, ] [[package]] -name = "nvidia-nvtx" -version = "13.0.85" +name = "nvidia-nvtx-cu12" +version = "12.6.77" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4", size = 148047, upload-time = "2025-09-04T08:29:01.761Z" }, - { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" }, + { url = "https://files.pythonhosted.org/packages/b9/93/80f8a520375af9d7ee44571a6544653a176e53c2b8ccce85b97b83c2491b/nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f44f8d86bb7d5629988d61c8d3ae61dddb2015dee142740536bc7481b022fe4b", size = 90549, upload-time = "2024-11-20T17:38:17.387Z" }, + { url = "https://files.pythonhosted.org/packages/2b/53/36e2fd6c7068997169b49ffc8c12d5af5e5ff209df6e1a2c4d373b3a638f/nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:adcaabb9d436c9761fca2b13959a2d237c5f9fd406c8e4b723c695409ff88059", size = 90539, upload-time = "2024-10-01T17:00:27.179Z" }, + { url = "https://files.pythonhosted.org/packages/56/9a/fff8376f8e3d084cd1530e1ef7b879bb7d6d265620c95c1b322725c694f4/nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b90bed3df379fa79afbd21be8e04a0314336b8ae16768b58f2d34cb1d04cd7d2", size = 89276, upload-time = "2024-11-20T17:38:27.621Z" }, + { url = "https://files.pythonhosted.org/packages/9e/4e/0d0c945463719429b7bd21dece907ad0bde437a2ff12b9b12fee94722ab0/nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:6574241a3ec5fdc9334353ab8c479fe75841dbe8f4532a8fc97ce63503330ba1", size = 89265, upload-time = "2024-10-01T17:00:38.172Z" }, ] [[package]] @@ -867,8 +873,10 @@ { name = "pyyaml" }, { name = "scikit-learn" }, { name = "timm" }, - { name = "torch" }, - { name = "torchvision" }, + { name = "torch", version = "2.14.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, + { name = "torch", version = "2.14.0+cu126", source = { registry = "https://download.pytorch.org/whl/cu126" }, marker = "sys_platform == 'linux'" }, + { name = "torchvision", version = "0.29.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, + { name = "torchvision", version = "0.29.0+cu126", source = { registry = "https://download.pytorch.org/whl/cu126" }, marker = "sys_platform == 'linux'" }, ] [package.dev-dependencies] @@ -886,8 +894,10 @@ { name = "pyyaml", specifier = ">=6.0" }, { name = "scikit-learn", specifier = ">=1.5" }, { name = "timm", specifier = ">=1.0" }, - { name = "torch", specifier = ">=2.5" }, - { name = "torchvision", specifier = ">=0.20" }, + { name = "torch", marker = "sys_platform != 'linux'", specifier = ">=2.5" }, + { name = "torch", marker = "sys_platform == 'linux'", specifier = ">=2.5", index = "https://download.pytorch.org/whl/cu126" }, + { name = "torchvision", marker = "sys_platform != 'linux'", specifier = ">=0.20" }, + { name = "torchvision", marker = "sys_platform == 'linux'", specifier = ">=0.20", index = "https://download.pytorch.org/whl/cu126" }, ] [package.metadata.requires-dev] @@ -904,8 +914,10 @@ { name = "huggingface-hub" }, { name = "pyyaml" }, { name = "safetensors" }, - { name = "torch" }, - { name = "torchvision" }, + { name = "torch", version = "2.14.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, + { name = "torch", version = "2.14.0+cu126", source = { registry = "https://download.pytorch.org/whl/cu126" }, marker = "sys_platform == 'linux'" }, + { name = "torchvision", version = "0.29.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" }, + { name = "torchvision", version = "0.29.0+cu126", source = { registry = "https://download.pytorch.org/whl/cu126" }, marker = "sys_platform == 'linux'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0e/ab/49041213990e48edcb6013463b18ef86d8d3c5d6ac5b408802c480443025/timm-1.0.30.tar.gz", hash = "sha256:1444a3af7e17fdba5440ac2bcf6e4b728da96dd62289ed9e1bcba31a545f5935", size = 2534738, upload-time = "2026-09-22T20:46:13.9Z" } wheels = [ @@ -916,46 +928,87 @@ name = "torch" version = "2.14.0" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "sys_platform != 'linux'", +] dependencies = [ - { name = "cuda-bindings", marker = "sys_platform == 'linux'" }, - { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, { name = "filelock" }, { name = "fsspec" }, { name = "jinja2" }, { name = "networkx" }, - { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" }, { name = "setuptools" }, { name = "sympy" }, - { name = "triton", marker = "sys_platform == 'linux'" }, { name = "typing-extensions" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/17/76/bb4770f56cf6d8971671dbcbb7493e5a6a15ad2825f4e359b02c27c38297/torch-2.14.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:c1f844f1c750e87df4b68bc3afbc0e2b0c7ef19d7b8f666e48bdcf6a0c4f0056", size = 127303200, upload-time = "2026-09-02T13:43:20.311Z" }, - { url = "https://files.pythonhosted.org/packages/fe/7b/ec44bacf2c8886b85ba4ca2285e8b09f2dff5d9c99e6a031326954082cb5/torch-2.14.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:ada340e62591d06a2bcc2d68170f20f45f0b0665d372dc510a8ed7eb3b1d609a", size = 454010251, upload-time = "2026-09-02T13:44:24.927Z" }, - { url = "https://files.pythonhosted.org/packages/15/71/49399acd41f750a906c686bd23c08a2001ccb8dd25f2971003c2ed89c1dd/torch-2.14.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:fecffb58f51fd643d213acd68da21cc3fc19bea05a3bc64b4ee55128f47a4963", size = 554620488, upload-time = "2026-09-02T13:44:49.442Z" }, { url = "https://files.pythonhosted.org/packages/be/16/9489b137112040f9911d7527e452854f21cc4e499ce0da79864e6a7451a7/torch-2.14.0-cp312-cp312-win_amd64.whl", hash = "sha256:cad84f41bbdf3dcf333ce394aeeaf25237c4d94fd6b659ba5eb813c829978823", size = 124114011, upload-time = "2026-09-02T13:43:55.034Z" }, ] [[package]] +name = "torch" +version = "2.14.0+cu126" +source = { registry = "https://download.pytorch.org/whl/cu126" } +resolution-markers = [ + "sys_platform == 'linux'", +] +dependencies = [ + { name = "cuda-bindings" }, + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvrtc", "nvtx"] }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "nvidia-cudnn-cu12" }, + { name = "nvidia-cusparselt-cu12" }, + { name = "nvidia-nccl-cu12" }, + { name = "nvidia-nvjitlink-cu12" }, + { name = "nvidia-nvshmem-cu12" }, + { name = "setuptools" }, + { name = "sympy" }, + { name = "triton" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://download-r2.pytorch.org/whl/cu126/torch-2.14.0%2Bcu126-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:b5723fa3dddd9cf548273eae0503ef9c47d0cce475edbf648ccab62224021fc7", upload-time = "2026-09-02T18:37:13Z" }, + { url = "https://download-r2.pytorch.org/whl/cu126/torch-2.14.0%2Bcu126-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:e9922441cb2ed269742a46246840c53634ca963f8b0a38d30182c28b4e8388da", upload-time = "2026-09-02T18:37:31Z" }, +] + +[[package]] name = "torchvision" version = "0.29.0" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "sys_platform != 'linux'", +] dependencies = [ { name = "numpy" }, { name = "pillow" }, - { name = "torch" }, + { name = "torch", version = "2.14.0", source = { registry = "https://pypi.org/simple" } }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/a7/d2/9daad500db2ab880eca0f0a569afd73e1afdd49f22c94560c85546dd87c2/torchvision-0.29.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:994687b818cac0e6d34cb407a41458e2c3262cb7db927465841c9522ba7ebb92", size = 1813740, upload-time = "2026-09-02T13:47:03.933Z" }, - { url = "https://files.pythonhosted.org/packages/41/14/702cd035fab1b8dfe944d06827cc4628c3b18e5d7ddaa676bf72c9be1c2c/torchvision-0.29.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:b7a736eaa6b2e22476c95ceb62986d195f3e600cdd7d196d7329aa2dfc951994", size = 7586112, upload-time = "2026-09-02T13:47:05.708Z" }, - { url = "https://files.pythonhosted.org/packages/27/d7/3cf1992414fd712af4865ae5f8f5758d9003c1ed82624908d54f7afb4342/torchvision-0.29.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:0e631a2f8b24732672d35224d2574ff89a78d56bef5364ef6094a625fa5f1771", size = 7430773, upload-time = "2026-09-02T13:47:07.178Z" }, { url = "https://files.pythonhosted.org/packages/ef/a8/59152b945c09840f582c40e38aa4128777625af7f2883ce2d77133f96a46/torchvision-0.29.0-cp312-cp312-win_amd64.whl", hash = "sha256:0e3b99c62f7f095153887330c2c60c9ebc2da84dfbe08eea564a41ec361629cf", size = 1376791, upload-time = "2026-09-02T13:47:11.283Z" }, ] [[package]] +name = "torchvision" +version = "0.29.0+cu126" +source = { registry = "https://download.pytorch.org/whl/cu126" } +resolution-markers = [ + "sys_platform == 'linux'", +] +dependencies = [ + { name = "numpy" }, + { name = "pillow" }, + { name = "torch", version = "2.14.0+cu126", source = { registry = "https://download.pytorch.org/whl/cu126" } }, +] +wheels = [ + { url = "https://download-r2.pytorch.org/whl/cu126/torchvision-0.29.0%2Bcu126-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:a4557ec86c5fc97b4d9f9477f20e2a387644701d05a1794c566621b2d413d807", upload-time = "2026-09-02T00:27:55Z" }, + { url = "https://download-r2.pytorch.org/whl/cu126/torchvision-0.29.0%2Bcu126-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:ec7dffc3d1defeefda5e9b45c39d5346eb6a1fc753f11920ec8ae3866d0f436c", upload-time = "2026-09-02T00:27:55Z" }, +] + +[[package]] name = "tqdm" version = "4.70.1" source = { registry = "https://pypi.org/simple" }