native single-worker kubeflow integration

documentation/OPTIONS.es.md

@@ -4,6 +4,21 @@
 
 Esta guía ofrece un desglose amigable de las opciones de línea de comandos disponibles en el script `train.py` de SimpleTuner. Estas opciones ofrecen un alto grado de personalización, permitiéndote entrenar tu modelo para ajustarlo a tus requisitos específicos.
 
+### Opciones del servidor Kubeflow para una sola GPU
+
+La integración opcional con Kubeflow delega la admisión de GPU a Kueue y crea Workers de un solo uso mediante Kubeflow Trainer. Instálalo con `pip install 'simpletuner[kubernetes]'` y ejecuta el Servidor dentro del clúster:
+
+```bash
+simpletuner server --mode trainer --kubeflow \
+  --kubeflow-namespace training \
+  --kubeflow-runtime simpletuner-worker \
+  --kubeflow-queue gpu-training \
+  --kubeflow-worker-image registry.example.com/simpletuner:latest \
+  --kubeflow-orchestrator-url http://simpletuner-server.training.svc:8001
+```
+
+`--kubeflow-poll-interval` controla la reconciliación del TrainJob y tiene como valor predeterminado cinco segundos. El administrador del clúster suministra el TrainingRuntime, la cola de Kueue, los permisos de ServiceAccount, la configuración de extracción de imágenes y el almacenamiento compartido de modelos/datos.
+
 ### Formato del archivo de configuración JSON
 
 El nombre de archivo JSON esperado es `config.json` y los nombres de clave son los mismos que los `--argumentos` de abajo. El prefijo `--` no es obligatorio en el archivo JSON, pero también puede dejarse.

documentation/OPTIONS.hi.md

@@ -4,6 +4,21 @@
 
 यह गाइड SimpleTuner के `train.py` स्क्रिप्ट में उपलब्ध command‑line विकल्पों का user‑friendly विवरण देती है। ये विकल्प उच्च स्तर का customization देते हैं, जिससे आप मॉडल को अपनी आवश्यकताओं के अनुसार ट्रेन कर सकते हैं।
 
+### Kubeflow single‑GPU server विकल्प
+
+वैकल्पिक Kubeflow integration GPU admission को Kueue को सौंप देता है और Kubeflow Trainer के माध्यम से एक‑बार उपयोग होने वाले Workers बनाता है। `pip install 'simpletuner[kubernetes]'` से install करें, फिर Server को cluster के अंदर चलाएँ:
+
+```bash
+simpletuner server --mode trainer --kubeflow \
+  --kubeflow-namespace training \
+  --kubeflow-runtime simpletuner-worker \
+  --kubeflow-queue gpu-training \
+  --kubeflow-worker-image registry.example.com/simpletuner:latest \
+  --kubeflow-orchestrator-url http://simpletuner-server.training.svc:8001
+```
+
+`--kubeflow-poll-interval` TrainJob reconciliation को नियंत्रित करता है और इसका default पाँच सेकंड है। Cluster administrator TrainingRuntime, Kueue queue, ServiceAccount permissions, image pull configuration और shared model/data storage प्रदान करता है।
+
 ### JSON Configuration file format
 
 अपेक्षित JSON फ़ाइल‑नाम `config.json` है और key नाम नीचे दिए `--arguments` जैसे ही हैं। JSON फ़ाइल में अग्रणी `--` आवश्यक नहीं है, लेकिन चाहें तो रख सकते हैं।

documentation/OPTIONS.ja.md

@@ -4,6 +4,21 @@
 
 このガイドでは、SimpleTuner の `train.py` スクリプトで利用できるコマンドラインオプションを分かりやすく説明します。これらのオプションにより高度なカスタマイズが可能になり、要件に合わせてモデルを学習できます。
 
+### Kubeflow シングル GPU サーバーオプション
+
+オプションの Kubeflow 統合は GPU アドミッションを Kueue に委任し、Kubeflow Trainer を通じて一回限りの Worker を作成します。`pip install 'simpletuner[kubernetes]'` でインストールし、クラスター内でサーバーを起動してください：
+
+```bash
+simpletuner server --mode trainer --kubeflow \
+  --kubeflow-namespace training \
+  --kubeflow-runtime simpletuner-worker \
+  --kubeflow-queue gpu-training \
+  --kubeflow-worker-image registry.example.com/simpletuner:latest \
+  --kubeflow-orchestrator-url http://simpletuner-server.training.svc:8001
+```
+
+`--kubeflow-poll-interval` は TrainJob の調整間隔を制御し、デフォルトは 5 秒です。クラスター管理者は TrainingRuntime、Kueue キュー、ServiceAccount 権限、イメージプル設定、および共有モデル/データストレージを用意します。
+
 ### JSON 設定ファイル形式
 
 期待される JSON ファイル名は `config.json` で、キー名は以下の `--arguments` と同じです。JSON では先頭の `--` は不要ですが、残していても構いません。

documentation/OPTIONS.md

@@ -4,6 +4,21 @@
 
 This guide provides a user-friendly breakdown of the command-line options available in SimpleTuner's `train.py` script. These options offer a high degree of customization, allowing you to train your model to suit your specific requirements.
 
+### Kubeflow single-GPU server options
+
+The optional Kubeflow integration delegates GPU admission to Kueue and creates one-shot Workers through Kubeflow Trainer. Install it with `pip install 'simpletuner[kubernetes]'`, then run the Server in-cluster:
+
+```bash
+simpletuner server --mode trainer --kubeflow \
+  --kubeflow-namespace training \
+  --kubeflow-runtime simpletuner-worker \
+  --kubeflow-queue gpu-training \
+  --kubeflow-worker-image registry.example.com/simpletuner:latest \
+  --kubeflow-orchestrator-url http://simpletuner-server.training.svc:8001
+```
+
+`--kubeflow-poll-interval` controls TrainJob reconciliation and defaults to five seconds. The cluster administrator supplies the TrainingRuntime, Kueue queue, ServiceAccount permissions, image pull configuration, and shared model/data storage.
+
 ### JSON Configuration file format
 
 The JSON filename expected is `config.json` and the key names are the same as the below `--arguments`. The leading `--` is not required for the JSON file, but it can be left in as well.

documentation/OPTIONS.pt-BR.md

@@ -4,6 +4,21 @@
 
 Este guia fornece uma explicacao amigavel das opcoes de linha de comando disponiveis no script `train.py` do SimpleTuner. Essas opcoes oferecem alto grau de customizacao, permitindo treinar seu modelo conforme suas necessidades.
 
+### Opções do servidor Kubeflow para GPU única
+
+A integração opcional com o Kubeflow delega a admissão de GPU ao Kueue e cria Workers de uso único por meio do Kubeflow Trainer. Instale com `pip install 'simpletuner[kubernetes]'` e execute o Servidor dentro do cluster:
+
+```bash
+simpletuner server --mode trainer --kubeflow \
+  --kubeflow-namespace training \
+  --kubeflow-runtime simpletuner-worker \
+  --kubeflow-queue gpu-training \
+  --kubeflow-worker-image registry.example.com/simpletuner:latest \
+  --kubeflow-orchestrator-url http://simpletuner-server.training.svc:8001
+```
+
+`--kubeflow-poll-interval` controla a reconciliação do TrainJob e tem padrão de cinco segundos. O administrador do cluster fornece o TrainingRuntime, a fila do Kueue, as permissões de ServiceAccount, a configuração de pull de imagem e o armazenamento compartilhado de modelos/dados.
+
 ### Formato do arquivo de configuracao JSON
 
 O nome esperado do arquivo JSON e `config.json` e os nomes das chaves sao os mesmos dos `--argumentos` abaixo. O prefixo `--` nao e necessario no arquivo JSON, mas pode ser mantido.

documentation/OPTIONS.zh.md

@@ -4,6 +4,21 @@
 
 本指南以易读方式介绍 SimpleTuner `train.py` 中可用的命令行选项。这些选项提供高度可定制性，使你能按需求训练模型。
 
+### Kubeflow 单 GPU 服务器选项
+
+可选的 Kubeflow 集成将 GPU 准入委托给 Kueue，并通过 Kubeflow Trainer 创建一次性 Worker。使用 `pip install 'simpletuner[kubernetes]'` 安装，然后在集群内运行服务器：
+
+```bash
+simpletuner server --mode trainer --kubeflow \
+  --kubeflow-namespace training \
+  --kubeflow-runtime simpletuner-worker \
+  --kubeflow-queue gpu-training \
+  --kubeflow-worker-image registry.example.com/simpletuner:latest \
+  --kubeflow-orchestrator-url http://simpletuner-server.training.svc:8001
+```
+
+`--kubeflow-poll-interval` 控制 TrainJob 调谐频率，默认为五秒。集群管理员负责提供 TrainingRuntime、Kueue 队列、ServiceAccount 权限、镜像拉取配置以及共享模型/数据存储。
+
 ### JSON 配置文件格式
 
 默认 JSON 文件名为 `config.json`，键名与下方 `--arguments` 一致。JSON 中不需要前导 `--`，但也可以保留。

setup.py

@@ -386,6 +386,7 @@ def _collect_package_files(*directories: str):
     "state-mysql": ["aiomysql>=0.2.0"],
     "state-redis": ["redis>=5.0.0"],
     "state-all": ["asyncpg>=0.29.0", "aiomysql>=0.2.0", "redis>=5.0.0"],
+    "kubernetes": ["kubernetes>=36.0.3,<37.0.0"],
     # All non-platform extras combined
     "all": [
         "pillow-jxl-plugin>=1.3.1",

simpletuner/cli/__init__.py

@@ -262,6 +262,39 @@ def _add_server_parser(subparsers):
         "-e",
         help="Environment to auto-start training with",
     )
+    server_parser.add_argument(
+        "--kubeflow",
+        action="store_true",
+        help="Enable one-shot single-GPU Workers through Kubeflow Trainer and Kueue",
+    )
+    server_parser.add_argument(
+        "--kubeflow-namespace",
+        default="default",
+        help="Kubernetes namespace for TrainJobs (default: default)",
+    )
+    server_parser.add_argument(
+        "--kubeflow-runtime",
+        default="simpletuner-worker",
+        help="Namespace-scoped TrainingRuntime name (default: simpletuner-worker)",
+    )
+    server_parser.add_argument(
+        "--kubeflow-queue",
+        help="Kueue LocalQueue used for TrainJob admission",
+    )
+    server_parser.add_argument(
+        "--kubeflow-worker-image",
+        help="Container image that runs the SimpleTuner Worker",
+    )
+    server_parser.add_argument(
+        "--kubeflow-orchestrator-url",
+        help="Server URL reachable from Worker Pods",
+    )
+    server_parser.add_argument(
+        "--kubeflow-poll-interval",
+        type=float,
+        default=5.0,
+        help="TrainJob reconciliation interval in seconds (default: 5)",
+    )
 
 
 def _add_shutdown_parser(subparsers):
@@ -381,9 +414,9 @@ def _add_jobs_parser(subparsers):
     submit_parser.add_argument(
         "--target",
         "-t",
-        choices=["local", "worker", "auto"],
+        choices=["local", "worker", "kubeflow", "auto"],
         default="auto",
-        help="Execution target: local (this machine), worker (remote), auto (prefer worker if available)",
+        help="Execution target: local, worker, kubeflow, or auto",
     )
 
     # cancel

simpletuner/cli/server.py

@@ -128,6 +128,49 @@ def _ensure_server_state_dir() -> None:
     os.environ["SIMPLETUNER_STATE_DIR"] = str(get_local_state_dir())
 
 
+def _configure_kubeflow_environment(args) -> Optional[str]:
+    """Validate Kubeflow CLI arguments and export the integration contract.
+
+    Args:
+        args: Parsed server command arguments.
+
+    Returns:
+        An error message when configuration is invalid, otherwise None.
+    """
+    if not getattr(args, "kubeflow", False):
+        return None
+    if getattr(args, "mode", "unified") == "callback":
+        return "Kubeflow scheduling requires trainer or unified server mode"
+
+    required = {
+        "--kubeflow-namespace": getattr(args, "kubeflow_namespace", None),
+        "--kubeflow-runtime": getattr(args, "kubeflow_runtime", None),
+        "--kubeflow-queue": getattr(args, "kubeflow_queue", None),
+        "--kubeflow-worker-image": getattr(args, "kubeflow_worker_image", None),
+        "--kubeflow-orchestrator-url": getattr(args, "kubeflow_orchestrator_url", None),
+    }
+    missing = [flag for flag, value in required.items() if not value]
+    if missing:
+        return f"Kubeflow scheduling requires: {', '.join(missing)}"
+
+    poll_interval = getattr(args, "kubeflow_poll_interval", 5.0)
+    if poll_interval <= 0:
+        return "--kubeflow-poll-interval must be greater than zero"
+
+    os.environ.update(
+        {
+            "SIMPLETUNER_KUBEFLOW_ENABLED": "true",
+            "SIMPLETUNER_KUBEFLOW_NAMESPACE": str(required["--kubeflow-namespace"]),
+            "SIMPLETUNER_KUBEFLOW_RUNTIME": str(required["--kubeflow-runtime"]),
+            "SIMPLETUNER_KUBEFLOW_QUEUE": str(required["--kubeflow-queue"]),
+            "SIMPLETUNER_KUBEFLOW_WORKER_IMAGE": str(required["--kubeflow-worker-image"]),
+            "SIMPLETUNER_KUBEFLOW_ORCHESTRATOR_URL": str(required["--kubeflow-orchestrator-url"]),
+            "SIMPLETUNER_KUBEFLOW_POLL_INTERVAL": str(poll_interval),
+        }
+    )
+    return None
+
+
 def cmd_server(args) -> int:
     """Handle server command."""
     host = getattr(args, "host", "0.0.0.0")
@@ -140,6 +183,11 @@ def cmd_server(args) -> int:
     ssl_no_verify = getattr(args, "ssl_no_verify", False)
     env = getattr(args, "env", None)
 
+    kubeflow_error = _configure_kubeflow_environment(args)
+    if kubeflow_error:
+        print(f"Error: {kubeflow_error}")
+        return 1
+
     _ensure_server_state_dir()
 
     if not os.environ.get("SIMPLETUNER_CONFIG_DIR"):

simpletuner/helpers/publishing/providers/s3.py

@@ -41,6 +41,22 @@ def __init__(self, config: dict[str, Any], *, provider_type: str = "s3"):
         self.use_ssl = bool(config.get("use_ssl", True))
         self._session = boto3.session.Session(**{k: v for k, v in session_kwargs.items() if v is not None})
         self._client = self._session.client("s3", endpoint_url=self.endpoint_url, use_ssl=self.use_ssl)
+        self.request_headers = {
+            str(name): str(value) for name, value in (config.get("request_headers") or {}).items()
+        }
+        self.force_single_part = bool(config.get("force_single_part", False))
+        if self.request_headers:
+            self._client.meta.events.register("before-sign.s3", self._inject_request_headers)
+
+    def _inject_request_headers(self, request: Any, **_: Any) -> None:
+        """Attach configured headers before an S3-compatible request is signed.
+
+        Args:
+            request: Botocore request being prepared for signing.
+            **_: Unused botocore event metadata.
+        """
+        for name, value in self.request_headers.items():
+            request.headers[name] = value
 
     def _build_uri(self, key_prefix: str) -> str:
         if self.public_base_url:
@@ -127,7 +143,11 @@ def publish(
         for file_path in files:
             relative_key = file_path.relative_to(path).as_posix() if path.is_dir() else file_path.name
             destination_key = "/".join([part for part in (destination_root, relative_key) if part])
-            self._client.upload_file(str(file_path), self.bucket, destination_key)
+            if self.force_single_part:
+                with file_path.open("rb") as body:
+                    self._client.put_object(Bucket=self.bucket, Key=destination_key, Body=body)
+            else:
+                self._client.upload_file(str(file_path), self.bucket, destination_key)
             last_key = destination_key
 
         assert last_key is not None  # for mypy/pylint; guarded by files check

simpletuner/helpers/training/trainer.py

@@ -610,6 +610,42 @@ def _init_publishing_manager(self):
             logger.error("Failed to initialise publishing providers: %s", exc)
             self.publishing_manager = None
 
+    def _publish_final_artifacts(self) -> List[Any]:
+        """Publish the completed output directory through configured providers.
+
+        Returns:
+            Publishing results produced by all successful providers.
+
+        Raises:
+            FileNotFoundError: If publishing is enabled but the output directory is missing.
+            RuntimeError: If a required publishing provider reports no artifacts.
+        """
+        manager = self.publishing_manager
+        if manager is None or not getattr(manager, "configured", False):
+            return []
+
+        output_dir = Path(self.config.output_dir)
+        if not output_dir.exists():
+            raise FileNotFoundError(f"Final publishing directory does not exist: {output_dir}")
+
+        results = manager.publish(
+            output_dir,
+            artifact_name=output_dir.name,
+            metadata={
+                "job_id": self.job_id,
+                "global_step": self.state.get("global_step", 0),
+                "epoch": self.state.get("current_epoch", 0),
+                "artifact_stage": "final",
+            },
+        )
+        required = any(
+            bool(getattr(provider, "config", {}).get("required", False))
+            for provider in getattr(manager, "providers", [])
+        )
+        if required and not results:
+            raise RuntimeError("Required final artifact publishing did not succeed")
+        return results
+
     def _enable_dynamo_dynamic_output_capture(self) -> None:
         try:
             import torch._dynamo as torch_dynamo
@@ -6999,6 +7035,12 @@ def _upload_final_model():
                 self._finish_hub_uploads()
             else:
                 self._run_post_upload_script(local_path=self.config.output_dir, remote_path=None)
+            if (
+                self.accelerator.is_main_process
+                and os.environ.get("SIMPLETUNER_PUBLISH_FINAL_ARTIFACTS", "").lower()
+                in {"1", "true", "yes", "on"}
+            ):
+                self._publish_final_artifacts()
             # Mark model_save as completed
             event = lifecycle_stage_event(
                 key="model_save",