HydroJobSche（Hydro Job Scheduler）

AI/ML
HydroJobSche（Hydro Job Scheduler）

HydroJobSche（Hydro Job Scheduler）

軽量なPyTorchトレーニング用ジョブスケジューラ。operator（Conv2d/Linear等）単位でフォワード／バックワードにフックを挿入し、Unixシグナル（SIGUSR1/SIGUSR2）で実行時にトレーニングを一時停止・再開できる設計。既存の学習ループをほとんど改修せずに組み込め、GPUリソースのプリエンプションやクラスタ運用でのきめ細かな制御を実現することを目指す。Python 3.7+/PyTorch 1.8+が前提。

https://github.com/ChenQiaoling00/HydroJobSche

概要

HydroJobSche は「実行時にトレーニングを細粒度で一時停止／再開」することに特化した、PyTorch向けの軽量ジョブスケジューラです。モデルの各オペレータ（たとえば Conv2d や Linear といったリーフモジュール）に対して自動的に前向き／後向きのフックを登録し、Unix シグナルを受け取ったときに学習処理を安全に止めたり再開したりできます。既存のトレーニングコードに侵襲せずにラップするだけで使えること、実行時オーバーヘッドを最小限にする設計、そして GPU リソースを短時間でプリエンプトするクラスタ運用シナリオでの利用が想定されています。

GitHub

リポジトリの統計情報

  • スター数: 9
  • フォーク数: 0
  • ウォッチャー数: 9
  • コミット数: 2
  • ファイル数: 7
  • メインの言語: Python

主な特徴

  • operator（モジュール）単位での細粒度な一時停止／再開を実現
  • Unix シグナル（SIGUSR1 / SIGUSR2）を使ったプロセス間制御
  • 既存モデルをラップする非侵襲的な導入（トレーニングループの大幅な変更不要）
  • モデルの葉モジュールへ再帰的に自動でフックを登録する自動化

技術的なポイント

HydroJobSche のコアは「フック」と「シグナルハンドラ」の組み合わせです。PyTorch のモジュールごとに forward/backward のフックを登録することで、演算の直前や直後に制御ポイントを作り、そこで現在の実行をブロック（待機）させることができます。プロセスは Unix シグナルで制御され、SIGUSR1 などのユーザー定義シグナルを受け取るとハンドラが状態フラグを立て、フック側はそのフラグを参照して処理を停止（たとえば threading.Event による wait）します。逆に再開用のシグナルでフラグがクリアされ、待機中のフックが解除されて処理を続行します。

この仕組みは「コード改修を最小限にしたい」運用上の要求にマッチします。HydroJobSche はモデルをラップするだけで、内部で再帰的に .children() を辿りながら各リーフモジュールにフックを貼り付けるため、既存の学習ループ・最適化コードをほとんど触らずに導入できます。実装面では PyTorch の register_forward_hook / register_backward_hook（新しい API を用いる実装も考えられる）や autograd のタイミングを用いて演算ごとに制御点を設けます。

設計上の注意点としては、これは「永続的なチェックポイント保存による復旧」ではなく、プロセス内での一時停止・再開を扱う仕組みであることです。つまりプロセスが強制終了した場合やノードが再起動した場合には状態が失われるため、長時間のプリエンプションや障害対応と組み合わせるなら外部チェックポイント機構（optimizer/state_dict の定期保存）と併用する必要があります。また、フック登録数が増えるとわずかなオーバーヘッドになるため、短い演算単位で頻繁にブロッキングすると性能に影響する可能性があります。GPUでの待機は CPU 側でブロックしている間、GPU カーネルの完了を待つ必要があり、同期ポイントが増えるとスループット低下を招くため、演算の粒度と待機ポリシーは調整が必要です。

実装の拡張ポイントとしては、シグナル以外に gRPC や REST API を用いた外部制御、チェックポイント連携、分散学習時の全ノード同期や部分停止サポート、そしてより洗練されたプリエンプションポリシー（優先度・最小停止時間など）の追加が考えられます。運用面では、ジョブスケジューラ（Kubernetes/SLURM等）と組み合わせてシグナル送出を自動化することで、GPU を効率的に再割当てするユースケースに適します。

プロジェクトの構成

主要なファイルとディレクトリ：

  • .DS_Store: file
  • HydroJobSche: dir
  • README.md: file
  • client.py: file
  • requirements.txt: file

…他 2 ファイル

まとめ

PyTorchトレーニングの一時停止／再開を非侵襲で実現する実装。クラスタ運用向けの軽量ソリューション。

リポジトリ情報：

READMEの抜粋：

HydroJobSche

HydroJobSche (Hydro Job Scheduler) 是一个支持在运行时暂停和恢复 PyTorch 深度学习训练任务的轻量级调度框架。通过在 operator 级别插入钩子（hooks），实现细粒度的训练控制，适用于 GPU 资源的抢占式调度场景。

🌟 核心特性

  • 细粒度控制：在每个 PyTorch operator（Conv2d、Linear 等）级别实现暂停/恢复
  • 信号驱动：使用 Unix 信号（SIGUSR1/SIGUSR2）进行进程间通信
  • 非侵入式设计：只需简单包装模型，无需修改训练代码
  • 自动钩子注册：递归地为模型的所有叶子模块注册前向/反向传播钩子
  • 轻量级实现：最小化运行时开销，适合生产环境

📦 安装

依赖要求

  • Python 3.7+
  • PyTorch 1.8+
  • CUDA（可选，用于 GPU 训练）

安装步骤

# 克隆仓库
cd HydroJobSche

# 安装依赖
pip install...