diff --git a/documentation/OPTIONS.es.md b/documentation/OPTIONS.es.md index ac141f309..87973d1fc 100644 --- a/documentation/OPTIONS.es.md +++ b/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. diff --git a/documentation/OPTIONS.hi.md b/documentation/OPTIONS.hi.md index 79f417525..f26fb7037 100644 --- a/documentation/OPTIONS.hi.md +++ b/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 फ़ाइल में अग्रणी `--` आवश्यक नहीं है, लेकिन चाहें तो रख सकते हैं। diff --git a/documentation/OPTIONS.ja.md b/documentation/OPTIONS.ja.md index cb67ab258..c41bd680b 100644 --- a/documentation/OPTIONS.ja.md +++ b/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 では先頭の `--` は不要ですが、残していても構いません。 diff --git a/documentation/OPTIONS.md b/documentation/OPTIONS.md index 12365de18..36f517e90 100644 --- a/documentation/OPTIONS.md +++ b/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. diff --git a/documentation/OPTIONS.pt-BR.md b/documentation/OPTIONS.pt-BR.md index c2c064547..e480ec51f 100644 --- a/documentation/OPTIONS.pt-BR.md +++ b/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. diff --git a/documentation/OPTIONS.zh.md b/documentation/OPTIONS.zh.md index 4a7d32dde..d4d54c68a 100644 --- a/documentation/OPTIONS.zh.md +++ b/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 中不需要前导 `--`，但也可以保留。 diff --git a/setup.py b/setup.py index fa0d8ac76..680d124ad 100644 --- a/setup.py +++ b/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", diff --git a/simpletuner/cli/__init__.py b/simpletuner/cli/__init__.py index 6c0f6882d..3d4877598 100644 --- a/simpletuner/cli/__init__.py +++ b/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 diff --git a/simpletuner/cli/server.py b/simpletuner/cli/server.py index 620546fea..fd4eecaa0 100644 --- a/simpletuner/cli/server.py +++ b/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"): diff --git a/simpletuner/helpers/publishing/providers/s3.py b/simpletuner/helpers/publishing/providers/s3.py index b6016de21..8c00cbad2 100644 --- a/simpletuner/helpers/publishing/providers/s3.py +++ b/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 diff --git a/simpletuner/helpers/training/trainer.py b/simpletuner/helpers/training/trainer.py index fb39ea0bf..2461c2c7d 100644 --- a/simpletuner/helpers/training/trainer.py +++ b/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", diff --git a/simpletuner/simpletuner_sdk/server/app.py b/simpletuner/simpletuner_sdk/server/app.py index 13bb01bf6..57de2c9be 100644 --- a/simpletuner/simpletuner_sdk/server/app.py +++ b/simpletuner/simpletuner_sdk/server/app.py @@ -411,35 +411,35 @@ async def lifespan(app: FastAPI): except Exception as e: logger.debug("Failed to configure rate limiters: %s", e) - # Initialize worker manager for GPU worker orchestration + # Initialize the shared Worker manager used by both ordinary Workers and + # one-shot Kubeflow Workers. worker_manager = None try: - from simpletuner.simpletuner_sdk.server.services.sse_manager import get_sse_manager - from simpletuner.simpletuner_sdk.server.services.worker_manager import WorkerManager - from simpletuner.simpletuner_sdk.server.services.worker_repository import WorkerRepository - - worker_repository = WorkerRepository() - - # Get job store for worker manager (uses cloud job store) from simpletuner.simpletuner_sdk.server.services.cloud.container import get_job_store + from simpletuner.simpletuner_sdk.server.services.sse_manager import get_sse_manager + from simpletuner.simpletuner_sdk.server.services.worker_manager import initialize_worker_manager + from simpletuner.simpletuner_sdk.server.services.worker_repository import get_worker_repository - job_store = get_job_store() - sse_manager = get_sse_manager() - - worker_manager = WorkerManager( - worker_repository=worker_repository, - job_store=job_store, - sse_manager=sse_manager, + worker_manager = await initialize_worker_manager( + worker_repository=get_worker_repository(), + job_store=get_job_store(), + sse_manager=get_sse_manager(), ) - - # Reconcile worker state from previous server run - await worker_manager.reconcile_on_startup() - - # Start background health check loop - await worker_manager.start() except Exception as e: logger.warning("Failed to start worker manager: %s", e) + # Kubeflow is opt-in. Invalid in-cluster configuration must fail startup + # instead of silently routing a GPU task through the local scheduler. + from simpletuner.simpletuner_sdk.server.services.kubeflow_job_service import ( + initialize_kubeflow_job_service, + ) + + kubeflow_job_service = await initialize_kubeflow_job_service() + if kubeflow_job_service is not None: + if worker_manager is None: + raise RuntimeError("Kubeflow scheduling requires an initialized WorkerManager") + logger.info("Kubeflow single-GPU job service started") + try: yield except asyncio.CancelledError: @@ -455,10 +455,22 @@ async def lifespan(app: FastAPI): except Exception as e: logger.warning("Error stopping background tasks: %s", e) - # Stop worker manager + # Stop Kubernetes reconciliation before the shared Worker manager. + try: + from simpletuner.simpletuner_sdk.server.services.kubeflow_job_service import ( + shutdown_kubeflow_job_service, + ) + + await shutdown_kubeflow_job_service() + logger.debug("Kubeflow job service stopped") + except Exception as e: + logger.warning("Error stopping Kubeflow job service: %s", e) + if worker_manager: try: - await worker_manager.stop() + from simpletuner.simpletuner_sdk.server.services.worker_manager import shutdown_worker_manager + + await shutdown_worker_manager() logger.debug("Worker manager stopped") except Exception as e: logger.warning("Error stopping worker manager: %s", e) diff --git a/simpletuner/simpletuner_sdk/server/models/worker.py b/simpletuner/simpletuner_sdk/server/models/worker.py index 389896d41..f3c069475 100644 --- a/simpletuner/simpletuner_sdk/server/models/worker.py +++ b/simpletuner/simpletuner_sdk/server/models/worker.py @@ -51,6 +51,19 @@ def persistent(self) -> bool: """Check if this is a persistent worker.""" return self.worker_type == WorkerType.PERSISTENT + @property + def is_job_bound(self) -> bool: + """Check whether Kubeflow provisioned this Worker for exactly one job. + + Returns: + True when the ephemeral Worker must retain its assigned job. + """ + return ( + self.worker_type == WorkerType.EPHEMERAL + and self.provider == "kubeflow" + and self.current_job_id is not None + ) + def to_dict(self) -> Dict[str, Any]: """Convert to dictionary for JSON serialization.""" return { diff --git a/simpletuner/simpletuner_sdk/server/routes/cloud/jobs.py b/simpletuner/simpletuner_sdk/server/routes/cloud/jobs.py index d065dbab6..e02235e06 100644 --- a/simpletuner/simpletuner_sdk/server/routes/cloud/jobs.py +++ b/simpletuner/simpletuner_sdk/server/routes/cloud/jobs.py @@ -199,6 +199,7 @@ async def cancel_job( CloudJobStatus.UPLOADING.value, CloudJobStatus.QUEUED.value, CloudJobStatus.RUNNING.value, + "starting", ] if job.status not in cancellable_states: raise JobStateError(job_id, job.status, cancellable_states) @@ -231,7 +232,19 @@ async def cancel_job( logger.error("Error cancelling cloud job %s: %s", job_id, exc) raise ProviderError(job.provider, f"Error cancelling job: {exc}", cause=exc) - if job.job_type == JobType.LOCAL: + is_kubeflow_job = job.job_type == JobType.LOCAL and job.provider == "kubeflow" + if is_kubeflow_job: + from ...services.kubeflow_job_service import get_kubeflow_job_service + + service = get_kubeflow_job_service() + if service is None: + raise ProviderError("kubeflow", "Kubeflow integration is not initialized") + try: + await service.cancel(job_id) + except Exception as exc: + logger.error("Error cancelling Kubeflow job %s: %s", job_id, exc) + raise ProviderError("kubeflow", f"Error cancelling job: {exc}", cause=exc) + elif job.job_type == JobType.LOCAL: try: from simpletuner.simpletuner_sdk import process_keeper @@ -311,7 +324,7 @@ async def cancel_job( logger.warning("Failed to broadcast SSE event: %s", exc) # Release GPUs and process pending jobs for local jobs AFTER broadcasting cancellation - if job.job_type == JobType.LOCAL: + if job.job_type == JobType.LOCAL and not is_kubeflow_job: try: from ...services.local_gpu_allocator import get_gpu_allocator diff --git a/simpletuner/simpletuner_sdk/server/routes/cloud/storage.py b/simpletuner/simpletuner_sdk/server/routes/cloud/storage.py index 37b229f45..bb94f0741 100644 --- a/simpletuner/simpletuner_sdk/server/routes/cloud/storage.py +++ b/simpletuner/simpletuner_sdk/server/routes/cloud/storage.py @@ -8,25 +8,84 @@ from __future__ import annotations +import hashlib import logging from pathlib import Path from typing import Any, Dict, Optional from fastapi import APIRouter, Depends, HTTPException, Request, status +from fastapi.responses import JSONResponse from ...services.cloud.auth import User, UserStore, get_optional_user +from ...services.kubeflow import KUBEFLOW_PROVIDER, LOCAL_UPLOAD_BUCKET from ._shared import get_client_ip, get_job_store, get_local_upload_dir logger = logging.getLogger(__name__) router = APIRouter(prefix="/storage", tags=["storage"]) +async def _record_kubeflow_upload( + store: Any, + authenticated_job: Any, + bucket: str, + key: str, +) -> None: + """Register one centrally received Kubeflow artifact. + + Args: + store: Unified job store. + authenticated_job: Job selected by the upload token. + bucket: S3-compatible bucket name. + key: Object key supplied by the publisher. + + Raises: + HTTPException: If a Kubeflow job writes outside its job-scoped prefix. + """ + if getattr(authenticated_job, "provider", None) != KUBEFLOW_PROVIDER: + return + + required_prefix = f"{authenticated_job.job_id}/" + if bucket != LOCAL_UPLOAD_BUCKET or not key.startswith(required_prefix): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Kubeflow artifacts must use the job-scoped output path", + ) + + metadata = dict(getattr(authenticated_job, "metadata", None) or {}) + upload_state = dict(metadata.get("artifact_upload") or {}) + received_files = list(upload_state.get("received_files") or []) + object_path = f"{bucket}/{key}" + if object_path not in received_files: + received_files.append(object_path) + upload_state.update( + { + "status": "receiving", + "received_files": received_files, + } + ) + metadata["artifact_upload"] = upload_state + await store.update_job( + authenticated_job.job_id, + { + "metadata": metadata, + "output_url": f"/api/cloud/storage/{LOCAL_UPLOAD_BUCKET}/{authenticated_job.job_id}", + }, + ) + @router.put("/{bucket}/{key:path}") -async def put_object(bucket: str, key: str, request: Request) -> Dict[str, Any]: +async def put_object(bucket: str, key: str, request: Request) -> JSONResponse: """S3-compatible PUT Object endpoint for receiving file uploads from the Cog. Requires authentication via per-job upload token. + + Args: + bucket: S3-compatible destination bucket. + key: Object key within the destination bucket. + request: Incoming upload request. + + Returns: + The legacy object metadata payload with an S3-compatible ETag header. """ store = get_job_store() client_ip = get_client_ip(request) @@ -77,12 +136,14 @@ async def put_object(bucket: str, key: str, request: Request) -> Dict[str, Any]: file_path.write_bytes(content) logger.info("Storage PUT: %s/%s (%d bytes)", bucket, key, len(content)) + await _record_kubeflow_upload(store, authenticated_job, bucket, key) - return { - "ETag": f'"{hash(content) & 0xFFFFFFFF:08x}"', - "Key": key, - "Bucket": bucket, - } + etag = hashlib.md5(content, usedforsecurity=False).hexdigest() + return JSONResponse( + status_code=status.HTTP_200_OK, + content={"ETag": f'"{etag}"', "Key": key, "Bucket": bucket}, + headers={"ETag": f'"{etag}"'}, + ) except HTTPException: raise diff --git a/simpletuner/simpletuner_sdk/server/routes/queue.py b/simpletuner/simpletuner_sdk/server/routes/queue.py index 0d95a54f4..1d3d5ee9c 100644 --- a/simpletuner/simpletuner_sdk/server/routes/queue.py +++ b/simpletuner/simpletuner_sdk/server/routes/queue.py @@ -355,18 +355,39 @@ async def cancel_queued_job( job_id: str, user: User = Depends(get_current_user), ) -> Dict[str, Any]: - """Cancel a job in the queue.""" + """Cancel a queued job while preserving provider-specific cleanup. + + Kubeflow jobs are represented in the unified job store instead of the + Scheduler queue. They must delete their TrainJob resources before the + request can be considered cancelled. + + Args: + job_id: Training job identifier. + user: Authenticated caller requesting cancellation. + + Returns: + A successful cancellation response. + + Raises: + HTTPException: If the job is not found, access is denied, or the + provider cannot cancel the requested job. + """ + from ..services.kubeflow_job_service import get_kubeflow_job_service + queue_store = QueueStore() entry = await queue_store.get_entry_by_job_id(job_id) + service = get_kubeflow_job_service() + kubeflow_job = await service.get_managed_job(job_id) if service else None - if not entry: + if not entry and kubeflow_job is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"Job not found in queue: {job_id}", ) # Check permission - is_own = entry.user_id == user.id + owner_id = kubeflow_job.user_id if kubeflow_job is not None else entry.user_id + is_own = owner_id == user.id can_cancel_all = user.has_permission("queue.cancel.all") can_cancel_own = user.has_permission("queue.cancel.own") @@ -381,6 +402,10 @@ async def cancel_queued_job( detail="Permission denied: queue.cancel.all", ) + if kubeflow_job is not None: + await service.cancel(job_id) + return {"success": True, "job_id": job_id} + scheduler = get_scheduler() success = await scheduler.cancel_job(job_id) @@ -579,9 +604,9 @@ class LocalJobSubmitRequest(BaseModel): no_wait: bool = Field(False, description="Reject immediately if GPUs unavailable") any_gpu: bool = Field(False, description="Use any available GPUs instead of configured device IDs") for_approval: bool = Field(False, description="Request approval to exceed org GPU quota") - target: Literal["local", "worker", "auto"] = Field( + target: Literal["local", "worker", "kubeflow", "auto"] = Field( "auto", - description="Execution target: 'local' (this machine), 'worker' (remote worker), 'auto' (prefer worker if available)", + description="Execution target: 'local', 'worker', 'kubeflow', or 'auto'", ) @@ -723,6 +748,7 @@ async def submit_local_job( Loads the config from disk and either: - Runs locally on this machine's GPUs (target='local') - Dispatches to a remote worker (target='worker') + - Creates a Kueue-admitted single-GPU TrainJob (target='kubeflow') - Auto-selects based on availability (target='auto', default) """ from pathlib import Path @@ -753,6 +779,31 @@ async def submit_local_job( error=f"Config '{request.config_name}' not found", ) + if request.target == "kubeflow": + from ..services.kubeflow_job_service import get_kubeflow_job_service + + service = get_kubeflow_job_service() + if service is None: + return LocalJobSubmitResponse( + success=False, + error="Kubeflow integration is not enabled on this server", + ) + try: + job = await service.submit( + config_name=request.config_name, + config=config, + user_id=user.id if user else None, + ) + except Exception as exc: + logger.exception("Failed to submit Kubeflow job") + return LocalJobSubmitResponse(success=False, error=str(exc)) + return LocalJobSubmitResponse( + success=True, + job_id=job.job_id, + status=job.status, + allocated_worker_id=(job.metadata or {}).get("worker_id"), + ) + # Always initialize worker repository to ensure consistent state # This fixes an issue where target=local would fail but target=auto would work worker_repo = None diff --git a/simpletuner/simpletuner_sdk/server/routes/workers.py b/simpletuner/simpletuner_sdk/server/routes/workers.py index bf43b6ee3..559a8fde3 100644 --- a/simpletuner/simpletuner_sdk/server/routes/workers.py +++ b/simpletuner/simpletuner_sdk/server/routes/workers.py @@ -8,7 +8,6 @@ from __future__ import annotations import asyncio -import hashlib import json import logging import secrets @@ -20,6 +19,7 @@ from pydantic import BaseModel, Field from ..services.cloud.auth import User, get_optional_user +from ..services.worker_credentials import generate_worker_token, hash_token logger = logging.getLogger(__name__) @@ -125,27 +125,6 @@ class TokenRotationResponse(BaseModel): # Helper Functions -def hash_token(token: str) -> str: - """Hash a worker token using SHA-256. - - Args: - token: The plaintext token - - Returns: - The hex digest of the token hash - """ - return hashlib.sha256(token.encode("utf-8")).hexdigest() - - -def generate_worker_token() -> str: - """Generate a secure worker token. - - Returns: - A URL-safe token string - """ - return secrets.token_urlsafe(32) - - async def validate_worker_token(token: str) -> Any: """Validate X-Worker-Token header and return worker. @@ -241,6 +220,7 @@ async def register_worker( Raises: HTTPException: 401 if token is invalid """ + from ..models.worker import WorkerStatus, WorkerType from ..services.worker_repository import get_worker_repository worker = await validate_worker_token(x_worker_token) @@ -257,21 +237,39 @@ async def register_worker( now = datetime.now(timezone.utc) - # Determine worker type - worker_type = "persistent" if request.persistent else "ephemeral" + if ( + worker.is_job_bound + and request.current_job_id is not None + and request.current_job_id != worker.current_job_id + ): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=( + f"Worker {worker.worker_id} is bound to job " + f"{worker.current_job_id}, not {request.current_job_id}" + ), + ) + + worker_type = ( + WorkerType.EPHEMERAL + if worker.is_job_bound + else WorkerType.PERSISTENT if request.persistent else WorkerType.EPHEMERAL + ) # Update worker with registration info updates = { "name": request.name, "gpu_info": request.gpu_info, - "status": "idle", - "worker_type": worker_type.upper(), + "status": WorkerStatus.CONNECTING if worker.is_job_bound else WorkerStatus.IDLE, + "worker_type": worker_type, "labels": request.labels, "last_heartbeat": now, "connected_at": now, } - if request.provider: + if worker.is_job_bound: + updates["current_job_id"] = worker.current_job_id + elif request.provider: updates["provider"] = request.provider await worker_repo.update_worker(worker.worker_id, updates) @@ -287,7 +285,10 @@ async def register_worker( job_repo = get_job_repository() job = await job_repo.get(request.current_job_id) - if job and job.status in ["running", "starting"]: + if worker.is_job_bound and job and job.status in ["pending", "queued"]: + # The SSE connection triggers the first and only dispatch. + pass + elif job and job.status in ["running", "starting"]: # Job is still active - worker should resume resume_job = { "job_id": job.job_id, @@ -342,6 +343,15 @@ async def worker_stream( queue = asyncio.Queue() worker_streams[worker_id] = queue + if worker.is_job_bound: + from ..services.worker_manager import get_worker_manager + + worker_manager = get_worker_manager() + if worker_manager is None: + logger.error("Cannot dispatch bound Worker %s: WorkerManager is not initialized", worker_id) + else: + await worker_manager.dispatch_bound_job(worker_id) + async def event_generator(): """Generate SSE events for the worker.""" try: @@ -411,15 +421,35 @@ async def worker_heartbeat( worker_repo = get_worker_repository() # Update worker with heartbeat info + if ( + worker.is_job_bound + and request.current_job_id is not None + and request.current_job_id != worker.current_job_id + ): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=( + f"Worker {worker.worker_id} is bound to job " + f"{worker.current_job_id}, not {request.current_job_id}" + ), + ) + now = datetime.now(timezone.utc) updates = { "status": request.status, "last_heartbeat": now, } + if worker.is_job_bound: + from ..models.worker import WorkerStatus + + updates["status"] = ( + WorkerStatus.BUSY if request.status == WorkerStatus.BUSY.value else WorkerStatus.CONNECTING + ) + updates["current_job_id"] = worker.current_job_id if worker.connected_at is None: updates["connected_at"] = now - if request.current_job_id: + if request.current_job_id and not worker.is_job_bound: updates["current_job_id"] = request.current_job_id await worker_repo.update_worker(worker.worker_id, updates) @@ -480,8 +510,31 @@ async def update_job_status( request.status, ) + completion_metadata = None + if request.status == "completed" and worker.is_job_bound: + from ..services.kubeflow_job_service import get_kubeflow_job_service + + service = get_kubeflow_job_service() + if service is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Kubeflow service is unavailable", + ) + if not service.artifacts_received(job): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="Central LoRA artifact upload has not completed", + ) + + completion_metadata = dict(job.metadata or {}) + upload_state = dict(completion_metadata.get("artifact_upload") or {}) + upload_state["status"] = "complete" + completion_metadata["artifact_upload"] = upload_state + # Update job status updates: Dict[str, Any] = {"status": request.status} + if completion_metadata is not None: + updates["metadata"] = completion_metadata if request.error: updates["error_message"] = request.error @@ -491,30 +544,47 @@ async def update_job_status( await job_repo.update(job_id, updates) - # Free worker when job completes + # Release an ordinary Worker, or delete a one-shot Kubeflow Worker and its + # TrainJob after the terminal status is safely persisted. if request.status in ["completed", "failed", "cancelled"]: from ..models.worker import WorkerStatus from ..services.worker_repository import get_worker_repository worker_repo = get_worker_repository() - await worker_repo.update_worker( - worker.worker_id, - { - "status": WorkerStatus.IDLE, - "current_job_id": None, - }, - ) - logger.info(f"Freed worker {worker.worker_id} after job {job_id} {request.status}") + if worker.is_job_bound: + await worker_repo.update_worker( + worker.worker_id, + { + "status": WorkerStatus.DRAINING, + "current_job_id": worker.current_job_id, + }, + ) + from ..services.kubeflow_job_service import get_kubeflow_job_service - # Try to dispatch any pending worker jobs - try: - from ..services.worker_manager import get_worker_manager + service = get_kubeflow_job_service() + if service is None: + logger.error("Kubeflow service unavailable while finalizing job %s", job_id) + else: + await service.finalize(job_id) + else: + await worker_repo.update_worker( + worker.worker_id, + { + "status": WorkerStatus.IDLE, + "current_job_id": None, + }, + ) + logger.info(f"Freed worker {worker.worker_id} after job {job_id} {request.status}") - worker_manager = get_worker_manager() - if worker_manager: - await worker_manager.dispatch_pending_jobs() - except Exception as exc: - logger.warning("Failed to dispatch pending jobs: %s", exc) + # Try to dispatch any pending worker jobs + try: + from ..services.worker_manager import get_worker_manager + + worker_manager = get_worker_manager() + if worker_manager: + await worker_manager.dispatch_pending_jobs() + except Exception as exc: + logger.warning("Failed to dispatch pending jobs: %s", exc) # Emit SSE event for UI updates try: diff --git a/simpletuner/simpletuner_sdk/server/services/cloud/job_logs.py b/simpletuner/simpletuner_sdk/server/services/cloud/job_logs.py index 9bcb5136e..6255c8cfa 100644 --- a/simpletuner/simpletuner_sdk/server/services/cloud/job_logs.py +++ b/simpletuner/simpletuner_sdk/server/services/cloud/job_logs.py @@ -88,6 +88,13 @@ async def fetch_job_logs(job: "UnifiedJob", max_bytes: int = 50000) -> str: Returns: Log content as a string """ + if job.provider == "kubeflow": + from ..kubeflow_job_service import get_kubeflow_job_service + + service = get_kubeflow_job_service() + if service is None: + return "(Kubeflow service is unavailable)" + return await service.get_logs(job) if job.job_type == JobType.CLOUD and job.provider: return await _fetch_cloud_logs(job) elif job.job_type == JobType.LOCAL: diff --git a/simpletuner/simpletuner_sdk/server/services/cloud/storage/job_repository.py b/simpletuner/simpletuner_sdk/server/services/cloud/storage/job_repository.py index f5d378887..987c08495 100644 --- a/simpletuner/simpletuner_sdk/server/services/cloud/storage/job_repository.py +++ b/simpletuner/simpletuner_sdk/server/services/cloud/storage/job_repository.py @@ -710,6 +710,7 @@ def _get(): """ SELECT * FROM jobs WHERE job_type = 'local' AND status IN ('pending', 'queued') + AND COALESCE(provider, '') != 'kubeflow' ORDER BY priority DESC, queued_at ASC """ ) @@ -731,6 +732,7 @@ def _get(): """ SELECT * FROM jobs WHERE job_type = 'local' AND status = 'running' + AND COALESCE(provider, '') != 'kubeflow' ORDER BY started_at ASC """ ) @@ -778,7 +780,13 @@ def _count(): conn = self._get_connection() try: cursor = conn.cursor() - cursor.execute("SELECT COUNT(*) as cnt FROM jobs WHERE job_type = 'local' AND status = 'running'") + cursor.execute( + """ + SELECT COUNT(*) as cnt FROM jobs + WHERE job_type = 'local' AND status = 'running' + AND COALESCE(provider, '') != 'kubeflow' + """ + ) return cursor.fetchone()["cnt"] finally: conn.close() diff --git a/simpletuner/simpletuner_sdk/server/services/kubeflow.py b/simpletuner/simpletuner_sdk/server/services/kubeflow.py new file mode 100644 index 000000000..9e29db7fc --- /dev/null +++ b/simpletuner/simpletuner_sdk/server/services/kubeflow.py @@ -0,0 +1,475 @@ +"""Kubernetes resources for one-shot, single-GPU SimpleTuner Workers.""" + +from __future__ import annotations + +import ast +import asyncio +import logging +import os +import re +from dataclasses import dataclass +from enum import Enum +from typing import Any, Mapping, Optional + +logger = logging.getLogger(__name__) + +KUBEFLOW_PROVIDER = "kubeflow" +LOCAL_UPLOAD_BUCKET = "outputs" +TRAINJOB_GROUP = "trainer.kubeflow.org" +TRAINJOB_VERSION = "v1alpha1" +TRAINJOB_PLURAL = "trainjobs" +GPU_RESOURCE_NAME = "nvidia.com/gpu" +SINGLE_GPU_QUANTITY = "1" +WORKER_TOKEN_KEY = "worker-token" + + +def _env_flag(value: Optional[str]) -> bool: + """Parse a conventional boolean environment value. + + Args: + value: Raw environment value. + + Returns: + Whether the value enables the feature. + """ + return bool(value and value.strip().lower() in {"1", "true", "yes", "on"}) + + +@dataclass(frozen=True, slots=True) +class KubeflowSettings: + """Configuration for namespace-scoped Kubeflow Worker provisioning.""" + + enabled: bool = False + namespace: str = "default" + runtime_name: str = "simpletuner-worker" + queue_name: Optional[str] = None + worker_image: Optional[str] = None + orchestrator_url: Optional[str] = None + poll_interval_seconds: float = 5.0 + + @classmethod + def from_env(cls, environ: Optional[Mapping[str, str]] = None) -> "KubeflowSettings": + """Load Kubeflow settings from environment variables. + + Args: + environ: Environment mapping, defaulting to the process environment. + + Returns: + Parsed and validated settings. + """ + values = os.environ if environ is None else environ + settings = cls( + enabled=_env_flag(values.get("SIMPLETUNER_KUBEFLOW_ENABLED")), + namespace=values.get("SIMPLETUNER_KUBEFLOW_NAMESPACE", "default"), + runtime_name=values.get("SIMPLETUNER_KUBEFLOW_RUNTIME", "simpletuner-worker"), + queue_name=values.get("SIMPLETUNER_KUBEFLOW_QUEUE"), + worker_image=values.get("SIMPLETUNER_KUBEFLOW_WORKER_IMAGE"), + orchestrator_url=values.get("SIMPLETUNER_KUBEFLOW_ORCHESTRATOR_URL"), + poll_interval_seconds=float(values.get("SIMPLETUNER_KUBEFLOW_POLL_INTERVAL", "5")), + ) + settings.validate() + return settings + + def validate(self) -> None: + """Validate settings required by an enabled integration. + + Raises: + ValueError: If an enabled integration is incomplete or invalid. + """ + if not self.enabled: + return + required = { + "queue_name": self.queue_name, + "worker_image": self.worker_image, + "orchestrator_url": self.orchestrator_url, + } + missing = [name for name, value in required.items() if not value] + if missing: + raise ValueError(f"Kubeflow configuration requires: {', '.join(missing)}") + if self.poll_interval_seconds <= 0: + raise ValueError("Kubeflow poll_interval_seconds must be greater than zero") + + +@dataclass(frozen=True, slots=True) +class KubeflowResources: + """Kubernetes resources allocated for one SimpleTuner job.""" + + namespace: str + trainjob_name: str + secret_name: str + trainjob_uid: Optional[str] = None + + def to_dict(self) -> dict[str, Optional[str]]: + """Serialize resource references into job metadata. + + Returns: + JSON-serializable resource references. + """ + return { + "namespace": self.namespace, + "trainjob_name": self.trainjob_name, + "secret_name": self.secret_name, + "trainjob_uid": self.trainjob_uid, + } + + @classmethod + def from_dict(cls, values: Mapping[str, Any]) -> "KubeflowResources": + """Restore resource references from job metadata. + + Args: + values: Kubernetes metadata stored with the job. + + Returns: + Resource reference object. + """ + return cls( + namespace=str(values["namespace"]), + trainjob_name=str(values["trainjob_name"]), + secret_name=str(values["secret_name"]), + trainjob_uid=values.get("trainjob_uid"), + ) + + +class KubeflowPhase(str, Enum): + """Infrastructure phase of a provisioned TrainJob.""" + + WAITING = "waiting_resource" + STARTING = "starting" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + MISSING = "missing" + + +def _resource_name(prefix: str, identifier: str) -> str: + """Build a deterministic DNS-compatible Kubernetes resource name. + + Args: + prefix: Resource-specific prefix. + identifier: SimpleTuner job identifier. + + Returns: + Kubernetes resource name no longer than 63 characters. + """ + normalized = re.sub(r"[^a-z0-9-]+", "-", identifier.lower()).strip("-") + return f"{prefix}-{normalized}"[:63].rstrip("-") + + +def _decode_log_payload(payload: Any) -> str: + """Normalize Kubernetes client log responses into plain UTF-8 text. + + Some Kubernetes client versions expose a response body as a stringified + bytes literal such as ``b"line\\n"``. This helper handles both raw bytes + and that representation without changing ordinary string logs. + + Args: + payload: Raw value returned by ``read_namespaced_pod_log``. + + Returns: + Decoded log text suitable for API responses and central archival. + """ + if isinstance(payload, bytes): + return payload.decode("utf-8", errors="replace") + + text = str(payload or "") + if len(text) >= 3 and text.startswith(("b'", 'b"')): + try: + decoded = ast.literal_eval(text) + except (SyntaxError, ValueError): + return text + if isinstance(decoded, bytes): + return decoded.decode("utf-8", errors="replace") + return text + + +class KubeflowWorkerProvisioner: + """Create and observe one-shot SimpleTuner TrainJobs.""" + + def __init__( + self, + settings: KubeflowSettings, + *, + core_api: Any = None, + custom_objects_api: Any = None, + ) -> None: + """Initialize the provisioner. + + Args: + settings: Validated Kubeflow settings. + core_api: Optional injected CoreV1Api for tests. + custom_objects_api: Optional injected CustomObjectsApi for tests. + """ + settings.validate() + self.settings = settings + if core_api is None or custom_objects_api is None: + core_api, custom_objects_api = self._load_incluster_clients() + self.core_api = core_api + self.custom_objects_api = custom_objects_api + + @staticmethod + def _load_incluster_clients() -> tuple[Any, Any]: + """Load official Kubernetes clients using the Pod ServiceAccount. + + Returns: + CoreV1Api and CustomObjectsApi instances. + + Raises: + RuntimeError: If the optional dependency is missing. + """ + try: + from kubernetes import client, config + except ImportError as exc: + raise RuntimeError( + "Kubeflow support requires the optional dependency: pip install 'simpletuner[kubernetes]'" + ) from exc + config.load_incluster_config() + return client.CoreV1Api(), client.CustomObjectsApi() + + def build_trainjob_manifest( + self, + *, + job_id: str, + worker_id: str, + secret_name: str, + ) -> dict[str, Any]: + """Build a TrainJob fixed to one Pod, process, and GPU. + + Args: + job_id: SimpleTuner job identifier. + worker_id: Pre-created Worker identifier. + secret_name: Secret containing the Worker token. + + Returns: + Kubeflow Trainer v2 TrainJob manifest. + """ + labels = { + "kueue.x-k8s.io/queue-name": str(self.settings.queue_name), + "simpletuner.ai/job-id": job_id, + "simpletuner.ai/worker-id": worker_id, + } + return { + "apiVersion": f"{TRAINJOB_GROUP}/{TRAINJOB_VERSION}", + "kind": "TrainJob", + "metadata": { + "name": _resource_name("simpletuner", job_id), + "namespace": self.settings.namespace, + "labels": labels, + }, + "spec": { + "suspend": True, + "runtimeRef": { + "name": self.settings.runtime_name, + "kind": "TrainingRuntime", + "apiGroup": TRAINJOB_GROUP, + }, + "trainer": { + "image": self.settings.worker_image, + "command": ["simpletuner", "worker"], + "numNodes": 1, + "numProcPerNode": 1, + "resourcesPerNode": { + "requests": {GPU_RESOURCE_NAME: SINGLE_GPU_QUANTITY}, + "limits": {GPU_RESOURCE_NAME: SINGLE_GPU_QUANTITY}, + }, + "env": [ + { + "name": "SIMPLETUNER_ORCHESTRATOR_URL", + "value": self.settings.orchestrator_url, + }, + { + "name": "SIMPLETUNER_WORKER_TOKEN", + "valueFrom": { + "secretKeyRef": { + "name": secret_name, + "key": WORKER_TOKEN_KEY, + } + }, + }, + {"name": "SIMPLETUNER_WORKER_NAME", "value": worker_id}, + {"name": "SIMPLETUNER_WORKER_PERSISTENT", "value": "false"}, + ], + }, + }, + } + + def build_secret_manifest( + self, + *, + job_id: str, + worker_token: str, + ) -> dict[str, Any]: + """Build the short-lived Worker authentication Secret. + + Args: + job_id: SimpleTuner job identifier. + worker_token: Plaintext Worker startup token. + + Returns: + Kubernetes Secret manifest. + """ + return { + "apiVersion": "v1", + "kind": "Secret", + "metadata": { + "name": _resource_name("simpletuner-worker", job_id), + "namespace": self.settings.namespace, + "labels": {"simpletuner.ai/job-id": job_id}, + }, + "type": "Opaque", + "stringData": {WORKER_TOKEN_KEY: worker_token}, + } + + async def create( + self, + *, + job_id: str, + worker_id: str, + worker_token: str, + ) -> KubeflowResources: + """Create the Worker Secret and single-GPU TrainJob. + + Args: + job_id: SimpleTuner job identifier. + worker_id: Bound Worker identifier. + worker_token: Plaintext Worker startup token. + + Returns: + Created Kubernetes resource references. + + Raises: + Exception: If either Kubernetes resource cannot be created. + """ + secret = self.build_secret_manifest(job_id=job_id, worker_token=worker_token) + secret_name = secret["metadata"]["name"] + trainjob = self.build_trainjob_manifest( + job_id=job_id, + worker_id=worker_id, + secret_name=secret_name, + ) + trainjob_name = trainjob["metadata"]["name"] + + await asyncio.to_thread( + self.core_api.create_namespaced_secret, + namespace=self.settings.namespace, + body=secret, + ) + try: + created = await asyncio.to_thread( + self.custom_objects_api.create_namespaced_custom_object, + group=TRAINJOB_GROUP, + version=TRAINJOB_VERSION, + namespace=self.settings.namespace, + plural=TRAINJOB_PLURAL, + body=trainjob, + ) + except Exception: + await asyncio.to_thread( + self.core_api.delete_namespaced_secret, + name=secret_name, + namespace=self.settings.namespace, + ) + raise + + return KubeflowResources( + namespace=self.settings.namespace, + trainjob_name=trainjob_name, + secret_name=secret_name, + trainjob_uid=created.get("metadata", {}).get("uid"), + ) + + async def get_phase(self, trainjob_name: str) -> KubeflowPhase: + """Read and normalize a TrainJob infrastructure phase. + + Args: + trainjob_name: Kubernetes TrainJob name. + + Returns: + Normalized infrastructure phase. + """ + try: + trainjob = await asyncio.to_thread( + self.custom_objects_api.get_namespaced_custom_object, + group=TRAINJOB_GROUP, + version=TRAINJOB_VERSION, + namespace=self.settings.namespace, + plural=TRAINJOB_PLURAL, + name=trainjob_name, + ) + except Exception as exc: + if getattr(exc, "status", None) == 404: + return KubeflowPhase.MISSING + raise + + conditions = trainjob.get("status", {}).get("conditions", []) + true_conditions = { + condition.get("type") + for condition in conditions + if str(condition.get("status", "")).lower() == "true" + } + if "Failed" in true_conditions: + return KubeflowPhase.FAILED + if "Complete" in true_conditions or "Succeeded" in true_conditions: + return KubeflowPhase.COMPLETED + if "Running" in true_conditions: + return KubeflowPhase.RUNNING + if trainjob.get("spec", {}).get("suspend", False): + return KubeflowPhase.WAITING + return KubeflowPhase.STARTING + + async def get_logs(self, trainjob_name: str) -> str: + """Read the current Worker Pod log for a TrainJob. + + Args: + trainjob_name: Kubernetes TrainJob name. + + Returns: + Combined Worker log text, or an empty string before Pod creation. + """ + pod_list = await asyncio.to_thread( + self.core_api.list_namespaced_pod, + namespace=self.settings.namespace, + label_selector=f"jobset.sigs.k8s.io/jobset-name={trainjob_name}", + ) + pods = list(getattr(pod_list, "items", None) or []) + if not pods: + return "" + pod_name = pods[0].metadata.name + logs = await asyncio.to_thread( + self.core_api.read_namespaced_pod_log, + name=pod_name, + namespace=self.settings.namespace, + timestamps=True, + ) + return _decode_log_payload(logs) + + async def delete(self, resources: KubeflowResources) -> None: + """Delete the TrainJob and Worker token Secret. + + Args: + resources: Kubernetes resources owned by the job. + """ + calls = ( + ( + self.custom_objects_api.delete_namespaced_custom_object, + { + "group": TRAINJOB_GROUP, + "version": TRAINJOB_VERSION, + "namespace": resources.namespace, + "plural": TRAINJOB_PLURAL, + "name": resources.trainjob_name, + }, + ), + ( + self.core_api.delete_namespaced_secret, + { + "name": resources.secret_name, + "namespace": resources.namespace, + }, + ), + ) + for operation, kwargs in calls: + try: + await asyncio.to_thread(operation, **kwargs) + except Exception as exc: + if getattr(exc, "status", None) != 404: + raise diff --git a/simpletuner/simpletuner_sdk/server/services/kubeflow_job_service.py b/simpletuner/simpletuner_sdk/server/services/kubeflow_job_service.py new file mode 100644 index 000000000..41595de9b --- /dev/null +++ b/simpletuner/simpletuner_sdk/server/services/kubeflow_job_service.py @@ -0,0 +1,460 @@ +"""Lifecycle service for Kubeflow-backed SimpleTuner jobs.""" + +from __future__ import annotations + +import asyncio +import copy +import logging +import secrets +import uuid +from datetime import datetime, timezone +from typing import Any, Optional +from urllib.parse import urlparse + +from ..models.worker import WorkerType +from .cloud.base import CloudJobStatus, JobType, UnifiedJob +from .kubeflow import ( + KUBEFLOW_PROVIDER, + LOCAL_UPLOAD_BUCKET, + KubeflowPhase, + KubeflowResources, + KubeflowSettings, + KubeflowWorkerProvisioner, +) +from .worker_credentials import create_worker_credentials + +logger = logging.getLogger(__name__) + + +class KubeflowJobService: + """Coordinate SimpleTuner jobs with one-shot Kubeflow Workers.""" + + def __init__( + self, + *, + settings: KubeflowSettings, + provisioner: KubeflowWorkerProvisioner, + job_store: Any, + worker_repository: Any, + ) -> None: + """Initialize the lifecycle service. + + Args: + settings: Validated Kubeflow configuration. + provisioner: Kubernetes resource provisioner. + job_store: Unified SimpleTuner job store. + worker_repository: Worker persistence repository. + """ + settings.validate() + self.settings = settings + self.provisioner = provisioner + self.job_store = job_store + self.worker_repository = worker_repository + self._task: Optional[asyncio.Task] = None + self._stop_event = asyncio.Event() + + async def submit( + self, + *, + config_name: str, + config: dict[str, Any], + user_id: Optional[int], + ) -> UnifiedJob: + """Submit one config as a queued single-GPU TrainJob. + + Args: + config_name: Stored SimpleTuner configuration name. + config: Fully materialized training configuration. + user_id: Submitting user identifier. + + Returns: + Persisted queued job. + """ + suffix = uuid.uuid4().hex[:12] + job_id = f"kjob-{suffix}" + worker_id = f"worker-{suffix}" + credentials = await create_worker_credentials( + self.worker_repository, + worker_id=worker_id, + name=worker_id, + user_id=user_id or 0, + worker_type=WorkerType.EPHEMERAL, + provider=KUBEFLOW_PROVIDER, + current_job_id=job_id, + ) + now = datetime.now(timezone.utc).isoformat() + upload_token = secrets.token_urlsafe(32) + job_config = self._with_local_publishing_config(config, job_id, upload_token) + job = UnifiedJob( + job_id=job_id, + job_type=JobType.LOCAL, + provider=KUBEFLOW_PROVIDER, + status=CloudJobStatus.QUEUED.value, + config_name=config_name, + created_at=now, + queued_at=now, + user_id=user_id, + num_processes=1, + upload_token=upload_token, + output_url=f"/api/cloud/storage/{LOCAL_UPLOAD_BUCKET}/{job_id}", + metadata={ + "config": job_config, + "target": KUBEFLOW_PROVIDER, + "worker_id": worker_id, + "provisioner": KUBEFLOW_PROVIDER, + "infrastructure_phase": KubeflowPhase.WAITING.value, + }, + ) + await self.job_store.add_job(job) + + try: + resources = await self.provisioner.create( + job_id=job_id, + worker_id=worker_id, + worker_token=credentials.token, + ) + except Exception as exc: + await self.worker_repository.delete_worker(worker_id) + await self.job_store.update_job( + job_id, + { + "status": CloudJobStatus.FAILED.value, + "completed_at": datetime.now(timezone.utc).isoformat(), + "error_message": f"Kubeflow TrainJob creation failed: {exc}", + }, + ) + raise + + job.metadata["kubernetes"] = resources.to_dict() + await self.job_store.update_job(job_id, {"metadata": job.metadata}) + return job + + def _with_local_publishing_config( + self, + config: dict[str, Any], + job_id: str, + upload_token: str, + ) -> dict[str, Any]: + """Append the Server's existing S3-compatible output destination. + + The resulting training configuration is consumed by the normal + SimpleTuner publishing path. The Worker remains unaware of Kubeflow. + + Args: + config: User training configuration. + job_id: Job identifier used as the remote object prefix. + upload_token: Per-job credential accepted by the Server. + + Returns: + A copied configuration containing the central publishing target. + + Raises: + ValueError: If an existing publishing configuration has an + unsupported representation. + """ + prepared = copy.deepcopy(config) + existing = prepared.get("publishing_config") + if existing in (None, "", [], {}, "None"): + publishing: list[dict[str, Any]] = [] + elif isinstance(existing, dict): + publishing = [existing] + elif isinstance(existing, list): + publishing = list(existing) + else: + raise ValueError("Kubeflow jobs require publishing_config to be a mapping or list") + + endpoint = f"{str(self.settings.orchestrator_url).rstrip('/')}/api/cloud/storage" + publishing.append( + { + "provider": "s3", + "bucket": LOCAL_UPLOAD_BUCKET, + "endpoint_url": endpoint, + "access_key": "local", + "secret_key": upload_token, + "base_path": job_id, + "use_ssl": urlparse(endpoint).scheme == "https", + "request_headers": {"X-SimpleTuner-Secret": upload_token}, + "force_single_part": True, + "required": True, + } + ) + prepared["publishing_config"] = publishing + return prepared + + + @staticmethod + def artifacts_received(job: UnifiedJob) -> bool: + """Check whether the Server has received a LoRA weight artifact. + + Args: + job: Kubeflow job whose upload metadata should be inspected. + + Returns: + True when at least one centrally registered safetensors file exists. + """ + metadata = getattr(job, "metadata", None) or {} + upload_state = metadata.get("artifact_upload") or {} + received_files = upload_state.get("received_files") or [] + return any( + str(object_path).lower().endswith(".safetensors") + for object_path in received_files + ) + + async def _archive_logs(self, job: UnifiedJob) -> dict[str, Any]: + """Persist the ephemeral Worker log before Kubernetes cleanup. + + Args: + job: Kubeflow job whose Worker Pod is still available. + + Returns: + Updated job metadata, including the archived log path when present. + """ + metadata = dict(job.metadata or {}) + try: + logs = await self.get_logs(job) + except Exception as exc: + logger.warning("Could not read Kubeflow logs for %s: %s", job.job_id, exc) + return metadata + if not isinstance(logs, str) or not logs: + return metadata + + from ..routes.cloud.helpers import get_local_upload_dir + + relative_path = f"{LOCAL_UPLOAD_BUCKET}/{job.job_id}/worker.log" + log_path = get_local_upload_dir() / relative_path + log_path.parent.mkdir(parents=True, exist_ok=True) + log_path.write_text(logs, encoding="utf-8") + + upload_state = dict(metadata.get("artifact_upload") or {}) + received_files = list(upload_state.get("received_files") or []) + if relative_path not in received_files: + received_files.append(relative_path) + upload_state["received_files"] = received_files + metadata["artifact_upload"] = upload_state + metadata["worker_log"] = relative_path + await self.job_store.update_job(job.job_id, {"metadata": metadata}) + return metadata + + async def get_logs(self, job: UnifiedJob) -> str: + """Read current Kubernetes logs through the Server adapter. + + Args: + job: Kubeflow job containing TrainJob resource metadata. + + Returns: + Current Worker Pod log text. + """ + resources = self._resources_for_job(job) + return await self.provisioner.get_logs(resources.trainjob_name) + + async def get_managed_job(self, job_id: str) -> Optional[UnifiedJob]: + """Return a job only when this service owns its runtime lifecycle. + + Args: + job_id: SimpleTuner job identifier. + + Returns: + The Kubeflow-backed job, or None for an unknown or differently + provisioned job. + """ + job = await self.job_store.get_job(job_id) + if job is None or job.provider != KUBEFLOW_PROVIDER: + return None + return job + + async def cancel(self, job_id: str) -> None: + """Cancel a Kubeflow job and release all ephemeral resources. + + Args: + job_id: SimpleTuner job identifier. + + Raises: + ValueError: If the job does not exist or lacks resource metadata. + """ + job = await self.get_managed_job(job_id) + if job is None: + raise ValueError(f"Kubeflow-managed job not found: {job_id}") + + worker_id = job.metadata.get("worker_id") + if worker_id: + from ..routes.workers import is_worker_connected, push_to_worker + + if is_worker_connected(worker_id): + await push_to_worker(worker_id, {"type": "job_cancel", "job_id": job_id}) + + resources = self._resources_for_job(job) + metadata = await self._archive_logs(job) + metadata["infrastructure_phase"] = CloudJobStatus.CANCELLED.value + await self.provisioner.delete(resources) + if worker_id: + await self.worker_repository.delete_worker(worker_id) + await self.job_store.update_job( + job_id, + { + "status": CloudJobStatus.CANCELLED.value, + "completed_at": datetime.now(timezone.utc).isoformat(), + "metadata": metadata, + }, + ) + + async def finalize(self, job_id: str) -> None: + """Release resources after a Worker reports a terminal state. + + Args: + job_id: Completed, failed, or cancelled job identifier. + """ + job = await self.job_store.get_job(job_id) + if job is None: + return + metadata = await self._archive_logs(job) + metadata["infrastructure_phase"] = job.status + await self.job_store.update_job(job_id, {"metadata": metadata}) + await self.provisioner.delete(self._resources_for_job(job)) + worker_id = job.metadata.get("worker_id") + if worker_id: + await self.worker_repository.delete_worker(worker_id) + + async def reconcile_once(self) -> None: + """Synchronize active TrainJob infrastructure state into job metadata.""" + jobs = await self.job_store.list_jobs(limit=1000, provider=KUBEFLOW_PROVIDER) + for job in jobs: + if job.is_terminal or not job.metadata.get("kubernetes"): + continue + resources = self._resources_for_job(job) + phase = await self.provisioner.get_phase(resources.trainjob_name) + metadata = dict(job.metadata) + metadata["infrastructure_phase"] = phase.value + updates: dict[str, Any] = {"metadata": metadata} + terminal = False + if phase in {KubeflowPhase.FAILED, KubeflowPhase.MISSING}: + terminal = True + updates.update( + { + "status": CloudJobStatus.FAILED.value, + "completed_at": datetime.now(timezone.utc).isoformat(), + "error_message": "Kubeflow TrainJob failed before Worker completion", + } + ) + elif phase == KubeflowPhase.COMPLETED: + terminal = True + terminal_status = ( + CloudJobStatus.COMPLETED.value + if self.artifacts_received(job) + else CloudJobStatus.FAILED.value + ) + updates.update( + { + "status": terminal_status, + "completed_at": datetime.now(timezone.utc).isoformat(), + } + ) + if terminal_status == CloudJobStatus.FAILED.value: + updates["error_message"] = ( + "Kubeflow Worker exited without confirmed central artifact upload" + ) + + if terminal: + archived_metadata = await self._archive_logs(job) + archived_metadata["infrastructure_phase"] = phase.value + updates["metadata"] = archived_metadata + await self.provisioner.delete(resources) + worker_id = job.metadata.get("worker_id") + if worker_id: + await self.worker_repository.delete_worker(worker_id) + await self.job_store.update_job(job.job_id, updates) + + async def start(self) -> None: + """Start periodic TrainJob reconciliation.""" + if self._task and not self._task.done(): + return + self._stop_event.clear() + self._task = asyncio.create_task(self._reconcile_loop()) + + async def stop(self) -> None: + """Stop periodic TrainJob reconciliation.""" + self._stop_event.set() + if self._task: + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + self._task = None + + async def _reconcile_loop(self) -> None: + """Run reconciliation until service shutdown.""" + while not self._stop_event.is_set(): + try: + await self.reconcile_once() + except Exception: + logger.exception("Kubeflow job reconciliation failed") + try: + await asyncio.wait_for( + self._stop_event.wait(), + timeout=self.settings.poll_interval_seconds, + ) + except asyncio.TimeoutError: + continue + + @staticmethod + def _resources_for_job(job: UnifiedJob) -> KubeflowResources: + """Load Kubernetes resource references from a job. + + Args: + job: Unified job containing Kubernetes metadata. + + Returns: + Kubernetes resource references. + + Raises: + ValueError: If resource metadata is absent. + """ + values = job.metadata.get("kubernetes") + if not values: + raise ValueError(f"Job {job.job_id} has no Kubernetes resource metadata") + return KubeflowResources.from_dict(values) + + +_kubeflow_job_service: Optional[KubeflowJobService] = None + + +def get_kubeflow_job_service() -> Optional[KubeflowJobService]: + """Return the initialized Kubeflow job service, if enabled.""" + return _kubeflow_job_service + + +async def initialize_kubeflow_job_service( + settings: Optional[KubeflowSettings] = None, +) -> Optional[KubeflowJobService]: + """Initialize the singleton Kubeflow job service when enabled. + + Args: + settings: Optional explicit settings for startup and tests. + + Returns: + Running service, or None when Kubeflow integration is disabled. + """ + global _kubeflow_job_service + resolved = settings or KubeflowSettings.from_env() + if not resolved.enabled: + return None + if _kubeflow_job_service is None: + from .cloud.container import get_job_store + from .worker_repository import get_worker_repository + + _kubeflow_job_service = KubeflowJobService( + settings=resolved, + provisioner=KubeflowWorkerProvisioner(resolved), + job_store=get_job_store(), + worker_repository=get_worker_repository(), + ) + await _kubeflow_job_service.start() + return _kubeflow_job_service + + +async def shutdown_kubeflow_job_service() -> None: + """Stop and clear the singleton Kubeflow job service.""" + global _kubeflow_job_service + if _kubeflow_job_service is not None: + await _kubeflow_job_service.stop() + _kubeflow_job_service = None diff --git a/simpletuner/simpletuner_sdk/server/services/worker_credentials.py b/simpletuner/simpletuner_sdk/server/services/worker_credentials.py new file mode 100644 index 000000000..4284eab1b --- /dev/null +++ b/simpletuner/simpletuner_sdk/server/services/worker_credentials.py @@ -0,0 +1,80 @@ +"""Credential creation shared by manual and provisioned GPU workers.""" + +from __future__ import annotations + +import hashlib +import secrets +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Optional + +from ..models.worker import Worker, WorkerStatus, WorkerType + + +@dataclass(frozen=True, slots=True) +class WorkerCredentials: + """A persisted Worker and its one-time plaintext token.""" + + worker: Worker + token: str + + +def hash_token(token: str) -> str: + """Hash a worker token using SHA-256. + + Args: + token: Plaintext Worker token. + + Returns: + Hexadecimal SHA-256 digest. + """ + return hashlib.sha256(token.encode("utf-8")).hexdigest() + + +def generate_worker_token() -> str: + """Generate a cryptographically secure Worker token. + + Returns: + URL-safe plaintext token. + """ + return secrets.token_urlsafe(32) + + +async def create_worker_credentials( + worker_repository, + *, + worker_id: str, + name: str, + user_id: int, + worker_type: WorkerType = WorkerType.EPHEMERAL, + provider: Optional[str] = None, + current_job_id: Optional[str] = None, +) -> WorkerCredentials: + """Create a Worker record and return its plaintext startup token. + + Args: + worker_repository: Repository used to persist the Worker. + worker_id: Stable Worker identifier. + name: Human-readable Worker name. + user_id: Owning SimpleTuner user identifier. + worker_type: Persistent or ephemeral Worker type. + provider: Infrastructure provider that owns the Worker. + current_job_id: Job permanently assigned to this Worker, if any. + + Returns: + Persisted Worker and its plaintext token. + """ + token = generate_worker_token() + worker = Worker( + worker_id=worker_id, + name=name, + worker_type=worker_type, + status=WorkerStatus.CONNECTING, + token_hash=hash_token(token), + user_id=user_id, + provider=provider, + current_job_id=current_job_id, + created_at=datetime.now(timezone.utc), + ) + await worker_repository.create_worker(worker) + return WorkerCredentials(worker=worker, token=token) diff --git a/simpletuner/simpletuner_sdk/server/services/worker_manager.py b/simpletuner/simpletuner_sdk/server/services/worker_manager.py index 34e41d0e0..77c67bb0d 100644 --- a/simpletuner/simpletuner_sdk/server/services/worker_manager.py +++ b/simpletuner/simpletuner_sdk/server/services/worker_manager.py @@ -100,6 +100,12 @@ async def _check_worker(self, worker, now: datetime): """Check individual worker health.""" from ..models.worker import WorkerStatus + # The Kubeflow service owns launch timeout, Pod failure, and cleanup for + # one-shot Workers. Generic recovery would otherwise requeue the job + # and delete the Worker while its TrainJob is still being admitted. + if worker.is_job_bound: + return + if worker.status == WorkerStatus.CONNECTING: time_since_created = now - worker.created_at if time_since_created > self.connecting_timeout: @@ -242,19 +248,56 @@ async def _find_available_worker(self, job): labels=required_labels, ) - async def _dispatch_job_to_worker(self, job, worker): - """Assign job to worker and push via SSE.""" + async def dispatch_bound_job(self, worker_id: str) -> bool: + """Dispatch the single job preassigned to a provisioned Worker. + + Args: + worker_id: Worker that just established its SSE stream. + + Returns: + True when the assigned job was pushed to that Worker. + """ + worker = await self.worker_repository.get_worker(worker_id) + if worker is None or not worker.is_job_bound: + return False + + job = await self.job_store.get_job(worker.current_job_id) + if ( + job is None + or getattr(job, "provider", None) != "kubeflow" + or job.status not in {"pending", "queued"} + ): + logger.warning( + "Bound worker %s has no dispatchable Kubeflow job %s", + worker_id, + worker.current_job_id, + ) + return False + + return await self._dispatch_job_to_worker(job, worker, preserve_binding=True) + + async def _dispatch_job_to_worker(self, job, worker, *, preserve_binding: bool = False) -> bool: + """Assign a job to a Worker and push it through SSE. + + Args: + job: Unified job to dispatch. + worker: Destination Worker. + preserve_binding: Keep the fixed job assignment if the push fails. + + Returns: + True when the event was pushed successfully. + """ from ..models.worker import WorkerStatus try: from ..routes.workers import is_worker_connected, push_to_worker except ImportError: logger.warning("Worker routes not available, cannot dispatch job") - return + return False if not is_worker_connected(worker.worker_id): logger.warning(f"Worker {worker.worker_id} not connected, skipping dispatch") - return + return False await self.worker_repository.update_worker( worker.worker_id, @@ -287,6 +330,7 @@ async def _dispatch_job_to_worker(self, job, worker): "job_id": job.job_id, "config": config, "dataloader": config.get("dataloader_config"), + "provider": getattr(job, "provider", None), "upload_endpoint": "/api/cloud/storage", "upload_token": getattr(job, "upload_token", None), "hf_token": job_metadata.get("hf_token"), @@ -295,26 +339,31 @@ async def _dispatch_job_to_worker(self, job, worker): if success: logger.info(f"Dispatched job {job.job_id} to worker {worker.worker_id}") - else: - logger.error(f"Failed to dispatch job {job.job_id} to worker {worker.worker_id}") - await self.worker_repository.update_worker( - worker.worker_id, - { - "status": WorkerStatus.IDLE, - "current_job_id": None, - }, - ) - # Reset job to pending, clear worker assignment from metadata - reset_metadata = dict(job.metadata or {}) + return True + + logger.error(f"Failed to dispatch job {job.job_id} to worker {worker.worker_id}") + await self.worker_repository.update_worker( + worker.worker_id, + { + "status": WorkerStatus.CONNECTING if preserve_binding else WorkerStatus.IDLE, + "current_job_id": worker.current_job_id if preserve_binding else None, + }, + ) + + # Bound jobs remain attached to their provisioned Worker so only the + # same Pod can retry the SSE connection. + reset_metadata = dict(job.metadata or {}) + if not preserve_binding: reset_metadata.pop("worker_id", None) - await self.job_store.update_job( - job.job_id, - { - "status": "pending", - "started_at": None, - "metadata": reset_metadata, - }, - ) + await self.job_store.update_job( + job.job_id, + { + "status": "queued" if preserve_binding else "pending", + "started_at": None, + "metadata": reset_metadata, + }, + ) + return False async def reconcile_on_startup(self): """Called on server startup to reconcile state.""" diff --git a/simpletuner/worker_agent.py b/simpletuner/worker_agent.py index 1be2ddf07..99219d8c0 100644 --- a/simpletuner/worker_agent.py +++ b/simpletuner/worker_agent.py @@ -15,6 +15,7 @@ import argparse import asyncio +import copy import json import logging import os @@ -119,9 +120,34 @@ def __init__(self, config: WorkerConfig): self.worker_id: Optional[str] = None self.current_job: Optional[Dict[str, Any]] = None self.shutdown_requested = False - self.training_process: Optional[subprocess.Popen] = None + self.training_process: Optional[Any] = None self._heartbeat_task: Optional[asyncio.Task] = None + @staticmethod + def _is_async_process(process: Any) -> bool: + """Return whether a process uses asyncio's subprocess interface. + + Args: + process: Training process owned by this Worker. + + Returns: + True when the process must be awaited through asyncio. + """ + return isinstance(process, asyncio.subprocess.Process) + + def _is_process_running(self, process: Any) -> bool: + """Return whether either supported subprocess type is still active. + + Args: + process: Training process owned by this Worker. + + Returns: + True while the process has not reached a terminal status. + """ + if self._is_async_process(process): + return process.returncode is None + return process.poll() is None + async def run(self): """Main entry point""" # Setup signal handlers @@ -155,7 +181,7 @@ def _handle_shutdown_signal(self): logger.info("Shutdown signal received") self.shutdown_requested = True - if self.training_process and self.training_process.poll() is None: + if self.training_process and self._is_process_running(self.training_process): logger.info("Sending SIGTERM to training process") self.training_process.terminate() @@ -258,18 +284,26 @@ async def _start_job(self, event: Dict[str, Any]): # Write config files to temp directory job_dir = Path(f"/tmp/simpletuner_job_{job_id}") - job_dir.mkdir(exist_ok=True) + job_dir.mkdir(parents=True, exist_ok=True) config_path = job_dir / "config.json" - dataloader_path = job_dir / "dataloader.yaml" - - # Import yaml here to avoid import at module level - import yaml - - with config_path.open("w", encoding="utf-8") as f: - json.dump(event["config"], f) - with open(dataloader_path, "w") as f: - yaml.dump(event["dataloader"], f) + is_kubeflow_job = event.get("provider") == "kubeflow" + job_config = copy.deepcopy(event["config"]) + dataloader = event.get("dataloader") + if is_kubeflow_job and dataloader is not None: + dataloader_path = job_dir / "dataloader.json" + with dataloader_path.open("w", encoding="utf-8") as handle: + json.dump(dataloader, handle) + job_config["data_backend_config"] = str(dataloader_path) + elif dataloader is not None: + dataloader_path = job_dir / "dataloader.yaml" + import yaml + + with dataloader_path.open("w", encoding="utf-8") as handle: + yaml.safe_dump(dataloader, handle) + + with config_path.open("w", encoding="utf-8") as handle: + json.dump(job_config, handle) # Set up environment env = os.environ.copy() @@ -283,6 +317,8 @@ async def _start_job(self, event: Dict[str, Any]): ) if event.get("hf_token"): env["HF_TOKEN"] = event["hf_token"] + if is_kubeflow_job: + env["SIMPLETUNER_PUBLISH_FINAL_ARTIFACTS"] = "true" # Report starting status await self._report_job_status("starting") @@ -296,33 +332,50 @@ async def _start_job(self, event: Dict[str, Any]): ] logger.info(f"Launching training: {' '.join(cmd)}") - self.training_process = subprocess.Popen( - cmd, - env=env, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - ) + if is_kubeflow_job: + self.training_process = await asyncio.create_subprocess_exec( + *cmd, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + else: + self.training_process = subprocess.Popen( + cmd, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) # Monitor in background asyncio.create_task(self._monitor_training(job_dir)) async def _monitor_training(self, job_dir: Path): - """Monitor training process and report status""" - process = self.training_process + """Monitor the training subprocess and report its terminal status. - await self._report_job_status("training") + Args: + job_dir: Per-job directory containing materialized configuration. + """ + process = self.training_process + if process is None: + raise RuntimeError("Training process is not initialized") - # Stream output (could send to server for log streaming) - while process.poll() is None: + if self._is_async_process(process): + await self._report_job_status("running") if process.stdout: - line = process.stdout.readline() - if line: - # Parse progress from output if possible - logger.info(f"[training] {line.rstrip()}") - await asyncio.sleep(0.1) - - exit_code = process.returncode + while line := await process.stdout.readline(): + logger.info("[training] %s", line.decode("utf-8", errors="replace").rstrip()) + exit_code = await process.wait() + else: + await self._report_job_status("training") + while process.poll() is None: + if process.stdout: + line = process.stdout.readline() + if line: + logger.info("[training] %s", line.rstrip()) + await asyncio.sleep(0.1) + exit_code = process.returncode logger.info(f"Training finished with exit code: {exit_code}") if exit_code == 0: @@ -343,16 +396,25 @@ async def _monitor_training(self, job_dir: Path): async def _stop_current_job(self): """Stop the current training job""" - if self.training_process and self.training_process.poll() is None: + process = self.training_process + if process and self._is_process_running(process): logger.info("Stopping current job") - self.training_process.terminate() + process.terminate() # Wait for graceful shutdown - try: - self.training_process.wait(timeout=30) - except subprocess.TimeoutExpired: - logger.warning("Training process did not stop, killing") - self.training_process.kill() + if self._is_async_process(process): + try: + await asyncio.wait_for(process.wait(), timeout=30) + except asyncio.TimeoutError: + logger.warning("Training process did not stop, killing") + process.kill() + await process.wait() + else: + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + logger.warning("Training process did not stop, killing") + process.kill() if self.current_job: await self._report_job_status("cancelled") diff --git a/tests/test_kubeflow_cancel_route.py b/tests/test_kubeflow_cancel_route.py new file mode 100644 index 000000000..6425192fe --- /dev/null +++ b/tests/test_kubeflow_cancel_route.py @@ -0,0 +1,70 @@ +"""Tests for cancelling Kubeflow-backed local jobs.""" + +from __future__ import annotations + +import unittest +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +from simpletuner.simpletuner_sdk.server.routes.cloud.jobs import cancel_job +from simpletuner.simpletuner_sdk.server.services.cloud.base import JobType, UnifiedJob + + +class KubeflowCancelRouteTestCase(unittest.IsolatedAsyncioTestCase): + """Test that cancellation delegates to the Kubernetes lifecycle owner.""" + + async def test_cancel_starting_job_skips_local_process_and_gpu_allocators(self) -> None: + """Verify a starting Kubeflow job remains cancellable by its resource owner.""" + job = UnifiedJob( + job_id="kjob-123", + job_type=JobType.LOCAL, + provider="kubeflow", + status="starting", + config_name="sdxl-lora", + created_at=datetime.now(timezone.utc).isoformat(), + metadata={"worker_id": "worker-123"}, + ) + store = AsyncMock() + store.get_job.return_value = job + service = AsyncMock() + request = MagicMock() + request.client.host = "10.0.0.1" + + with ( + patch( + "simpletuner.simpletuner_sdk.server.routes.cloud.jobs.get_job_store", + return_value=store, + ), + patch( + "simpletuner.simpletuner_sdk.server.services.kubeflow_job_service.get_kubeflow_job_service", + return_value=service, + ), + patch( + "simpletuner.simpletuner_sdk.process_keeper.terminate_process", + return_value=True, + ) as terminate_process, + patch( + "simpletuner.simpletuner_sdk.server.services.local_gpu_allocator.get_gpu_allocator", + ) as get_gpu_allocator, + patch( + "simpletuner.simpletuner_sdk.server.services.cloud.audit.audit_log", + new=AsyncMock(), + ), + patch( + "simpletuner.simpletuner_sdk.server.services.sse_manager.get_sse_manager", + ) as get_sse_manager, + patch( + "simpletuner.simpletuner_sdk.server.routes.cloud.jobs.emit_cloud_event", + ), + ): + get_sse_manager.return_value = AsyncMock() + result = await cancel_job("kjob-123", request, user=None) + + self.assertTrue(result["success"]) + service.cancel.assert_awaited_once_with("kjob-123") + terminate_process.assert_not_called() + get_gpu_allocator.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_kubeflow_cli.py b/tests/test_kubeflow_cli.py new file mode 100644 index 000000000..2d9c1e430 --- /dev/null +++ b/tests/test_kubeflow_cli.py @@ -0,0 +1,65 @@ +"""Tests for Kubeflow server command configuration.""" + +from __future__ import annotations + +import os +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from simpletuner.cli import create_parser +from simpletuner.cli.server import _configure_kubeflow_environment + + +class KubeflowServerCliTestCase(unittest.TestCase): + """Test explicit Kubeflow server flags and environment mapping.""" + + def test_parser_accepts_single_gpu_kubeflow_options(self) -> None: + """Verify the server command exposes all required integration values.""" + args = create_parser().parse_args( + [ + "server", + "--kubeflow", + "--kubeflow-queue", + "gpu-training", + "--kubeflow-worker-image", + "registry.example.com/simpletuner:4.5.0", + "--kubeflow-orchestrator-url", + "http://simpletuner.simpletuner.svc:8001", + ] + ) + + self.assertTrue(args.kubeflow) + self.assertEqual(args.kubeflow_namespace, "default") + self.assertEqual(args.kubeflow_runtime, "simpletuner-worker") + self.assertEqual(args.kubeflow_poll_interval, 5.0) + + def test_environment_mapping_requires_and_exports_cluster_values(self) -> None: + """Verify validated CLI values become the server integration contract.""" + args = SimpleNamespace( + kubeflow=True, + mode="trainer", + kubeflow_namespace="simpletuner", + kubeflow_runtime="simpletuner-worker", + kubeflow_queue="gpu-training", + kubeflow_worker_image="registry.example.com/simpletuner:4.5.0", + kubeflow_orchestrator_url="http://simpletuner.simpletuner.svc:8001", + kubeflow_poll_interval=3.0, + ) + + with patch.dict(os.environ, {}, clear=True): + error = _configure_kubeflow_environment(args) + + self.assertIsNone(error) + self.assertEqual(os.environ["SIMPLETUNER_KUBEFLOW_ENABLED"], "true") + self.assertEqual(os.environ["SIMPLETUNER_KUBEFLOW_QUEUE"], "gpu-training") + self.assertEqual( + os.environ["SIMPLETUNER_KUBEFLOW_WORKER_IMAGE"], + "registry.example.com/simpletuner:4.5.0", + ) + + def test_callback_only_server_rejects_kubeflow(self) -> None: + """Verify the scheduler cannot run without trainer APIs.""" + args = SimpleNamespace(kubeflow=True, mode="callback") + + self.assertIn("trainer", _configure_kubeflow_environment(args)) diff --git a/tests/test_kubeflow_job_service.py b/tests/test_kubeflow_job_service.py new file mode 100644 index 000000000..a639e3127 --- /dev/null +++ b/tests/test_kubeflow_job_service.py @@ -0,0 +1,262 @@ +"""Tests for Kubeflow-backed SimpleTuner jobs.""" + +from __future__ import annotations + +import tempfile +import unittest +from datetime import datetime, timezone +from pathlib import Path +from unittest.mock import AsyncMock, patch + +from simpletuner.simpletuner_sdk.server.services.cloud.base import CloudJobStatus, JobType, UnifiedJob +from simpletuner.simpletuner_sdk.server.services.cloud.job_logs import fetch_job_logs +from simpletuner.simpletuner_sdk.server.services.kubeflow import KubeflowPhase, KubeflowResources, KubeflowSettings +from simpletuner.simpletuner_sdk.server.services.kubeflow_job_service import KubeflowJobService + + +class KubeflowJobServiceTestCase(unittest.IsolatedAsyncioTestCase): + """Test job persistence and one-shot Worker provisioning.""" + + def setUp(self) -> None: + """Create service fixtures.""" + self.settings = KubeflowSettings( + enabled=True, + namespace="simpletuner", + runtime_name="simpletuner-worker", + queue_name="gpu-training", + worker_image="registry.example.com/simpletuner:4.5.0", + orchestrator_url="http://simpletuner.simpletuner.svc:8001", + ) + self.job_store = AsyncMock() + self.worker_repository = AsyncMock() + self.provisioner = AsyncMock() + self.provisioner.create.return_value = KubeflowResources( + namespace="simpletuner", + trainjob_name="simpletuner-kjob-123", + secret_name="simpletuner-worker-kjob-123", + trainjob_uid="uid-123", + ) + self.service = KubeflowJobService( + settings=self.settings, + provisioner=self.provisioner, + job_store=self.job_store, + worker_repository=self.worker_repository, + ) + + async def test_submit_creates_bound_worker_and_single_process_job(self) -> None: + """Verify submission binds one ephemeral Worker to one queued job.""" + config = {"model_family": "sdxl"} + with ( + patch( + "simpletuner.simpletuner_sdk.server.services.kubeflow_job_service.uuid.uuid4" + ) as uuid4, + patch( + "simpletuner.simpletuner_sdk.server.services.worker_credentials.generate_worker_token", + return_value="test-token", + ), + patch( + "simpletuner.simpletuner_sdk.server.services.kubeflow_job_service.secrets.token_urlsafe", + return_value="test-upload-token", + create=True, + ), + ): + uuid4.return_value.hex = "1234567890abcdef" + job = await self.service.submit( + config_name="sdxl-lora", + config=config, + user_id=7, + ) + + self.assertEqual(job.provider, "kubeflow") + self.assertEqual(job.status, CloudJobStatus.QUEUED.value) + self.assertEqual(job.num_processes, 1) + self.assertEqual(job.metadata["worker_id"], "worker-1234567890ab") + self.assertEqual(job.metadata["target"], "kubeflow") + self.assertEqual(job.upload_token, "test-upload-token") + publishing = job.metadata["config"]["publishing_config"] + self.assertEqual( + publishing[-1], + { + "provider": "s3", + "bucket": "outputs", + "endpoint_url": "http://simpletuner.simpletuner.svc:8001/api/cloud/storage", + "access_key": "local", + "secret_key": "test-upload-token", + "base_path": job.job_id, + "use_ssl": False, + "request_headers": {"X-SimpleTuner-Secret": "test-upload-token"}, + "force_single_part": True, + "required": True, + }, + ) + self.assertNotIn("publishing_config", config) + self.job_store.add_job.assert_awaited_once() + created_worker = self.worker_repository.create_worker.await_args.args[0] + self.assertEqual(created_worker.current_job_id, job.job_id) + self.assertEqual(created_worker.provider, "kubeflow") + self.provisioner.create.assert_awaited_once_with( + job_id=job.job_id, + worker_id=created_worker.worker_id, + worker_token="test-token", + ) + + async def test_fetch_logs_reads_active_worker_pod(self) -> None: + """Verify the existing logs endpoint polls Kubeflow through the Server.""" + job = UnifiedJob( + job_id="kjob-123", + job_type=JobType.LOCAL, + provider="kubeflow", + status=CloudJobStatus.RUNNING.value, + config_name="sdxl-lora", + created_at=datetime.now(timezone.utc).isoformat(), + metadata={ + "kubernetes": { + "namespace": "simpletuner", + "trainjob_name": "simpletuner-kjob-123", + "secret_name": "simpletuner-worker-kjob-123", + } + }, + ) + self.provisioner.get_logs.return_value = "training step 1/1" + + with patch( + "simpletuner.simpletuner_sdk.server.services.kubeflow_job_service.get_kubeflow_job_service", + return_value=self.service, + ): + logs = await fetch_job_logs(job) + + self.provisioner.get_logs.assert_awaited_once_with("simpletuner-kjob-123") + self.assertEqual(logs, "training step 1/1") + + async def test_cancel_deletes_trainjob_secret_and_worker(self) -> None: + """Verify cancellation releases Kubernetes and Worker resources.""" + job = await self.service.submit( + config_name="sdxl-lora", + config={"model_family": "sdxl"}, + user_id=7, + ) + self.job_store.get_job.return_value = job + + with ( + patch( + "simpletuner.simpletuner_sdk.server.routes.workers.is_worker_connected", + return_value=False, + ), + patch( + "simpletuner.simpletuner_sdk.server.routes.workers.push_to_worker", + new_callable=AsyncMock, + ), + ): + await self.service.cancel(job.job_id) + + self.provisioner.delete.assert_awaited_once() + self.worker_repository.delete_worker.assert_awaited_once_with(job.metadata["worker_id"]) + terminal_update = self.job_store.update_job.await_args_list[-1].args[1] + self.assertEqual(terminal_update["status"], CloudJobStatus.CANCELLED.value) + self.assertEqual( + terminal_update["metadata"]["infrastructure_phase"], + CloudJobStatus.CANCELLED.value, + ) + + async def test_finalize_archives_worker_log_before_deletion(self) -> None: + """Verify Pod logs reach central storage before resource cleanup.""" + job = UnifiedJob( + job_id="kjob-123", + job_type=JobType.LOCAL, + provider="kubeflow", + status=CloudJobStatus.COMPLETED.value, + config_name="sdxl-lora", + created_at=datetime.now(timezone.utc).isoformat(), + metadata={ + "worker_id": "worker-123", + "infrastructure_phase": KubeflowPhase.STARTING.value, + "kubernetes": { + "namespace": "simpletuner", + "trainjob_name": "simpletuner-kjob-123", + "secret_name": "simpletuner-worker-kjob-123", + }, + }, + ) + self.job_store.get_job.return_value = job + events = [] + self.provisioner.get_logs.side_effect = lambda *_: events.append("logs") or "training complete" + self.provisioner.delete.side_effect = lambda *_: events.append("delete") + + with tempfile.TemporaryDirectory() as tmpdir, patch( + "simpletuner.simpletuner_sdk.server.routes.cloud.helpers.get_local_upload_dir", + return_value=Path(tmpdir), + ): + await self.service.finalize(job.job_id) + archived_log = Path(tmpdir) / "outputs" / job.job_id / "worker.log" + self.assertEqual(archived_log.read_text(encoding="utf-8"), "training complete") + + self.assertEqual(events, ["logs", "delete"]) + metadata_update = self.job_store.update_job.await_args.args[1]["metadata"] + self.assertIn( + f"outputs/{job.job_id}/worker.log", + metadata_update["artifact_upload"]["received_files"], + ) + self.assertEqual( + metadata_update["infrastructure_phase"], + KubeflowPhase.COMPLETED.value, + ) + + async def test_reconcile_failed_trainjob_releases_bound_resources(self) -> None: + """Verify an infrastructure failure cannot leak a GPU Worker.""" + job = UnifiedJob( + job_id="kjob-123", + job_type=JobType.LOCAL, + provider="kubeflow", + status=CloudJobStatus.QUEUED.value, + config_name="sdxl-lora", + created_at=datetime.now(timezone.utc).isoformat(), + metadata={ + "worker_id": "worker-123", + "kubernetes": { + "namespace": "simpletuner", + "trainjob_name": "simpletuner-kjob-123", + "secret_name": "simpletuner-worker-kjob-123", + }, + }, + ) + self.job_store.list_jobs.return_value = [job] + self.provisioner.get_phase.return_value = KubeflowPhase.FAILED + + await self.service.reconcile_once() + + self.provisioner.delete.assert_awaited_once() + self.worker_repository.delete_worker.assert_awaited_once_with("worker-123") + updates = self.job_store.update_job.await_args.args[1] + self.assertEqual(updates["status"], CloudJobStatus.FAILED.value) + + + async def test_reconcile_completed_trainjob_without_artifact_fails(self) -> None: + """Verify infrastructure success cannot bypass artifact confirmation.""" + job = UnifiedJob( + job_id="kjob-123", + job_type=JobType.LOCAL, + provider="kubeflow", + status=CloudJobStatus.RUNNING.value, + config_name="sdxl-lora", + created_at=datetime.now(timezone.utc).isoformat(), + metadata={ + "worker_id": "worker-123", + "kubernetes": { + "namespace": "simpletuner", + "trainjob_name": "simpletuner-kjob-123", + "secret_name": "simpletuner-worker-kjob-123", + }, + }, + ) + self.job_store.list_jobs.return_value = [job] + self.provisioner.get_phase.return_value = KubeflowPhase.COMPLETED + + await self.service.reconcile_once() + + updates = self.job_store.update_job.await_args.args[1] + self.assertEqual(updates["status"], CloudJobStatus.FAILED.value) + self.assertIn("artifact", updates["error_message"].lower()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_kubeflow_provisioner.py b/tests/test_kubeflow_provisioner.py new file mode 100644 index 000000000..846f4fc4d --- /dev/null +++ b/tests/test_kubeflow_provisioner.py @@ -0,0 +1,169 @@ +"""Tests for single-GPU Kubeflow TrainJob provisioning.""" + +from __future__ import annotations + +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock + +from simpletuner.simpletuner_sdk.server.services.kubeflow import ( + KubeflowPhase, + KubeflowSettings, + KubeflowWorkerProvisioner, +) + + +class KubeflowSettingsTestCase(unittest.TestCase): + """Test Kubeflow server configuration validation.""" + + def test_enabled_settings_require_runtime_values(self) -> None: + """Verify enabled integration rejects incomplete configuration.""" + settings = KubeflowSettings(enabled=True) + + with self.assertRaisesRegex(ValueError, "queue_name"): + settings.validate() + + +class KubeflowWorkerProvisionerTestCase(unittest.IsolatedAsyncioTestCase): + """Test TrainJob and Secret lifecycle without a live cluster.""" + + def setUp(self) -> None: + """Create provisioner fixtures.""" + self.settings = KubeflowSettings( + enabled=True, + namespace="simpletuner", + runtime_name="simpletuner-worker", + queue_name="gpu-training", + worker_image="registry.example.com/simpletuner:4.5.0", + orchestrator_url="http://simpletuner.simpletuner.svc:8001", + ) + self.core_api = MagicMock() + self.custom_api = MagicMock() + self.provisioner = KubeflowWorkerProvisioner( + self.settings, + core_api=self.core_api, + custom_objects_api=self.custom_api, + ) + + def test_trainjob_requests_exactly_one_gpu(self) -> None: + """Verify every generated TrainJob is one node, one process, one GPU.""" + manifest = self.provisioner.build_trainjob_manifest( + job_id="kjob-123", + worker_id="worker-123", + secret_name="simpletuner-worker-kjob-123", + ) + + trainer = manifest["spec"]["trainer"] + self.assertEqual(manifest["apiVersion"], "trainer.kubeflow.org/v1alpha1") + self.assertTrue(manifest["spec"]["suspend"]) + self.assertEqual(trainer["numNodes"], 1) + self.assertEqual(trainer["numProcPerNode"], 1) + self.assertEqual(trainer["resourcesPerNode"]["requests"], {"nvidia.com/gpu": "1"}) + self.assertEqual(trainer["resourcesPerNode"]["limits"], {"nvidia.com/gpu": "1"}) + self.assertEqual( + manifest["metadata"]["labels"]["kueue.x-k8s.io/queue-name"], + "gpu-training", + ) + + def test_worker_token_is_loaded_from_secret(self) -> None: + """Verify the TrainJob never contains the plaintext worker token.""" + manifest = self.provisioner.build_trainjob_manifest( + job_id="kjob-123", + worker_id="worker-123", + secret_name="simpletuner-worker-kjob-123", + ) + + token_env = next( + item + for item in manifest["spec"]["trainer"]["env"] + if item["name"] == "SIMPLETUNER_WORKER_TOKEN" + ) + self.assertEqual( + token_env["valueFrom"]["secretKeyRef"], + {"name": "simpletuner-worker-kjob-123", "key": "worker-token"}, + ) + self.assertNotIn("test-token", str(manifest)) + + def test_worker_uses_standard_simpletuner_environment(self) -> None: + """Verify Kubernetes does not require a Kubeflow-aware Worker.""" + manifest = self.provisioner.build_trainjob_manifest( + job_id="kjob-123", + worker_id="worker-123", + secret_name="simpletuner-worker-kjob-123", + ) + + trainer = manifest["spec"]["trainer"] + env_names = {item["name"] for item in trainer["env"]} + + self.assertEqual(trainer["command"], ["simpletuner", "worker"]) + self.assertNotIn("SIMPLETUNER_WORKER_PROVIDER", env_names) + self.assertNotIn("SIMPLETUNER_WORKER_BOUND_JOB_ID", env_names) + + async def test_create_rolls_back_secret_when_trainjob_creation_fails(self) -> None: + """Verify a failed TrainJob creation does not leave a token Secret.""" + self.custom_api.create_namespaced_custom_object.side_effect = RuntimeError("create failed") + + with self.assertRaisesRegex(RuntimeError, "create failed"): + await self.provisioner.create( + job_id="kjob-123", + worker_id="worker-123", + worker_token="test-token", + ) + + self.core_api.create_namespaced_secret.assert_called_once() + self.core_api.delete_namespaced_secret.assert_called_once_with( + name="simpletuner-worker-kjob-123", + namespace="simpletuner", + ) + + async def test_get_phase_maps_complete_condition(self) -> None: + """Verify completed TrainJobs map to a terminal provisioning phase.""" + self.custom_api.get_namespaced_custom_object.return_value = { + "status": { + "conditions": [ + {"type": "Complete", "status": "True"}, + ] + } + } + + phase = await self.provisioner.get_phase("simpletuner-kjob-123") + + self.assertEqual(phase, KubeflowPhase.COMPLETED) + + + async def test_get_logs_reads_trainjob_pod(self) -> None: + """Verify Server-side polling reads the ephemeral Worker Pod log.""" + self.core_api.list_namespaced_pod.return_value = SimpleNamespace( + items=[SimpleNamespace(metadata=SimpleNamespace(name="worker-pod"))] + ) + self.core_api.read_namespaced_pod_log.return_value = b"training step 1/1" + + logs = await self.provisioner.get_logs("simpletuner-kjob-123") + + self.core_api.list_namespaced_pod.assert_called_once_with( + namespace="simpletuner", + label_selector="jobset.sigs.k8s.io/jobset-name=simpletuner-kjob-123", + ) + self.core_api.read_namespaced_pod_log.assert_called_once_with( + name="worker-pod", + namespace="simpletuner", + timestamps=True, + ) + self.assertEqual(logs, "training step 1/1") + + async def test_get_logs_decodes_stringified_bytes(self) -> None: + """Verify Kubernetes client byte representations become plain text.""" + self.core_api.list_namespaced_pod.return_value = SimpleNamespace( + items=[SimpleNamespace(metadata=SimpleNamespace(name="worker-pod"))] + ) + self.core_api.read_namespaced_pod_log.return_value = ( + 'b"training step 1/1\\ntraining complete"' + ) + + logs = await self.provisioner.get_logs("simpletuner-kjob-123") + + self.assertEqual(logs, "training step 1/1\ntraining complete") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_kubeflow_queue_cancel_route.py b/tests/test_kubeflow_queue_cancel_route.py new file mode 100644 index 000000000..7e361b457 --- /dev/null +++ b/tests/test_kubeflow_queue_cancel_route.py @@ -0,0 +1,141 @@ +"""Tests for cancelling Kubeflow jobs through the generic queue API.""" + +from __future__ import annotations + +import unittest +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +from simpletuner.simpletuner_sdk.server.routes.queue import cancel_queued_job +from simpletuner.simpletuner_sdk.server.services.cloud.base import JobType, UnifiedJob + + +class KubeflowQueueCancelRouteTestCase(unittest.IsolatedAsyncioTestCase): + """Verify the generic queue cancellation endpoint is provider-transparent.""" + + async def test_cancel_kubeflow_job_uses_service_owned_record(self) -> None: + """Route a Kubeflow job by the lifecycle service's authoritative store.""" + job = UnifiedJob( + job_id="kjob-123", + job_type=JobType.LOCAL, + provider="kubeflow", + status="running", + config_name="sdxl-lora", + created_at=datetime.now(timezone.utc).isoformat(), + user_id=7, + ) + job_store = MagicMock() + job_store.get_job = AsyncMock(return_value=None) + queue_store = MagicMock() + queue_store.get_entry_by_job_id = AsyncMock(return_value=None) + service = MagicMock() + service.get_managed_job = AsyncMock(return_value=job) + service.cancel = AsyncMock() + user = MagicMock(id=7) + user.has_permission.return_value = True + + with ( + patch( + "simpletuner.simpletuner_sdk.server.routes.cloud.dependencies.get_job_store", + return_value=job_store, + ), + patch( + "simpletuner.simpletuner_sdk.server.routes.queue.QueueStore", + return_value=queue_store, + ), + patch( + "simpletuner.simpletuner_sdk.server.services.kubeflow_job_service.get_kubeflow_job_service", + return_value=service, + ), + patch("simpletuner.simpletuner_sdk.server.routes.queue.get_scheduler") as get_scheduler, + ): + result = await cancel_queued_job("kjob-123", user=user) + + self.assertEqual(result, {"success": True, "job_id": "kjob-123"}) + service.cancel.assert_awaited_once_with("kjob-123") + get_scheduler.assert_not_called() + + async def test_cancel_kubeflow_job_ignores_stale_compatibility_queue(self) -> None: + """Avoid delegating a managed job to the legacy Scheduler path.""" + job = UnifiedJob( + job_id="kjob-789", + job_type=JobType.LOCAL, + provider="kubeflow", + status="running", + config_name="sdxl-lora", + created_at=datetime.now(timezone.utc).isoformat(), + user_id=7, + ) + entry = MagicMock(user_id=7, provider="local") + job_store = MagicMock() + job_store.get_job = AsyncMock(return_value=None) + queue_store = MagicMock() + queue_store.get_entry_by_job_id = AsyncMock(return_value=entry) + queue_store.mark_cancelled_by_job_id = AsyncMock(return_value=True) + service = MagicMock() + service.get_managed_job = AsyncMock(return_value=job) + service.cancel = AsyncMock() + user = MagicMock(id=7) + user.has_permission.return_value = True + + with ( + patch( + "simpletuner.simpletuner_sdk.server.routes.cloud.dependencies.get_job_store", + return_value=job_store, + ), + patch( + "simpletuner.simpletuner_sdk.server.routes.queue.QueueStore", + return_value=queue_store, + ), + patch( + "simpletuner.simpletuner_sdk.server.services.kubeflow_job_service.get_kubeflow_job_service", + return_value=service, + ), + patch("simpletuner.simpletuner_sdk.server.routes.queue.get_scheduler") as get_scheduler, + ): + result = await cancel_queued_job("kjob-789", user=user) + + self.assertEqual(result, {"success": True, "job_id": "kjob-789"}) + service.cancel.assert_awaited_once_with("kjob-789") + get_scheduler.assert_not_called() + + async def test_cancel_regular_queue_job_keeps_scheduler_path(self) -> None: + """Preserve the existing Scheduler cancellation behavior for normal jobs.""" + entry = MagicMock(user_id=7) + job_store = MagicMock() + job_store.get_job = AsyncMock(return_value=None) + queue_store = MagicMock() + queue_store.get_entry_by_job_id = AsyncMock(return_value=entry) + scheduler = MagicMock() + scheduler.cancel_job = AsyncMock(return_value=True) + service = MagicMock() + service.get_managed_job = AsyncMock(return_value=None) + user = MagicMock(id=7) + user.has_permission.return_value = True + + with ( + patch( + "simpletuner.simpletuner_sdk.server.routes.cloud.dependencies.get_job_store", + return_value=job_store, + ), + patch( + "simpletuner.simpletuner_sdk.server.routes.queue.QueueStore", + return_value=queue_store, + ), + patch( + "simpletuner.simpletuner_sdk.server.services.kubeflow_job_service.get_kubeflow_job_service", + return_value=service, + ), + patch( + "simpletuner.simpletuner_sdk.server.routes.queue.get_scheduler", + return_value=scheduler, + ), + ): + result = await cancel_queued_job("job-456", user=user) + + self.assertEqual(result, {"success": True, "job_id": "job-456"}) + scheduler.cancel_job.assert_awaited_once_with("job-456") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_kubeflow_queue_route.py b/tests/test_kubeflow_queue_route.py new file mode 100644 index 000000000..5f8891f02 --- /dev/null +++ b/tests/test_kubeflow_queue_route.py @@ -0,0 +1,67 @@ +"""Tests for explicit Kubeflow queue submission.""" + +from __future__ import annotations + +import unittest +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +from simpletuner.simpletuner_sdk.server.routes.queue import LocalJobSubmitRequest, submit_local_job +from simpletuner.simpletuner_sdk.server.services.cloud.base import JobType, UnifiedJob + + +class KubeflowQueueRouteTestCase(unittest.IsolatedAsyncioTestCase): + """Test that Kubeflow submissions bypass SimpleTuner GPU selection.""" + + async def test_submit_kubeflow_bypasses_local_and_worker_allocators(self) -> None: + """Verify Kueue is the only GPU admission path for this target.""" + config_store = MagicMock() + config_store.load_config.return_value = ({"model_family": "sdxl"}, MagicMock()) + defaults = MagicMock(configs_dir="/configs") + state_store = MagicMock() + state_store.load_defaults.return_value = defaults + service = AsyncMock() + service.submit.return_value = UnifiedJob( + job_id="kjob-123", + job_type=JobType.LOCAL, + provider="kubeflow", + status="queued", + config_name="sdxl-lora", + created_at=datetime.now(timezone.utc).isoformat(), + metadata={"worker_id": "worker-123"}, + ) + + with ( + patch( + "simpletuner.simpletuner_sdk.server.services.webui_state.WebUIStateStore", + return_value=state_store, + ), + patch( + "simpletuner.simpletuner_sdk.server.services.config_store.ConfigStore", + return_value=config_store, + ), + patch( + "simpletuner.simpletuner_sdk.server.services.kubeflow_job_service.get_kubeflow_job_service", + return_value=service, + ), + patch( + "simpletuner.simpletuner_sdk.server.services.training_service.start_training_job" + ) as start_local, + patch( + "simpletuner.simpletuner_sdk.server.services.worker_repository.get_worker_repository" + ) as get_worker_repository, + ): + result = await submit_local_job( + LocalJobSubmitRequest(config_name="sdxl-lora", target="kubeflow"), + user=None, + ) + + self.assertTrue(result.success) + self.assertEqual(result.status, "queued") + self.assertEqual(result.allocated_worker_id, "worker-123") + start_local.assert_not_called() + get_worker_repository.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_kubeflow_scheduler_isolation.py b/tests/test_kubeflow_scheduler_isolation.py new file mode 100644 index 000000000..8f5d24a54 --- /dev/null +++ b/tests/test_kubeflow_scheduler_isolation.py @@ -0,0 +1,62 @@ +"""Tests that Kubeflow jobs never enter the local GPU scheduler.""" + +from __future__ import annotations + +import tempfile +import unittest +from datetime import datetime, timezone +from pathlib import Path + +from simpletuner.simpletuner_sdk.server.services.cloud.base import CloudJobStatus, JobType, UnifiedJob +from simpletuner.simpletuner_sdk.server.services.cloud.storage.job_repository import JobRepository + + +class KubeflowSchedulerIsolationTestCase(unittest.IsolatedAsyncioTestCase): + """Verify Kubernetes owns GPU admission for Kubeflow jobs.""" + + def setUp(self) -> None: + """Create an isolated job repository.""" + self.temp_dir = tempfile.TemporaryDirectory() + self.repository = JobRepository(Path(self.temp_dir.name) / "jobs.db") + + def tearDown(self) -> None: + """Remove the isolated repository.""" + self.repository.reset_instance() + self.temp_dir.cleanup() + + async def test_kubeflow_jobs_are_excluded_from_local_gpu_queries(self) -> None: + """Verify local allocation cannot select or count a Kubeflow job.""" + now = datetime.now(timezone.utc).isoformat() + local_job = UnifiedJob( + job_id="local-job", + job_type=JobType.LOCAL, + provider="local", + status=CloudJobStatus.QUEUED.value, + config_name="local-config", + created_at=now, + queued_at=now, + ) + kubeflow_job = UnifiedJob( + job_id="kubeflow-job", + job_type=JobType.LOCAL, + provider="kubeflow", + status=CloudJobStatus.QUEUED.value, + config_name="kubeflow-config", + created_at=now, + queued_at=now, + ) + await self.repository.add(local_job) + await self.repository.add(kubeflow_job) + + pending_jobs = await self.repository.get_pending_local_jobs() + self.assertEqual([job.job_id for job in pending_jobs], ["local-job"]) + + await self.repository.mark_running(local_job.job_id) + await self.repository.mark_running(kubeflow_job.job_id) + running_jobs = await self.repository.get_running_local_jobs() + self.assertEqual([job.job_id for job in running_jobs], ["local-job"]) + self.assertEqual(await self.repository.count_running_local_jobs(), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_kubeflow_worker_routes.py b/tests/test_kubeflow_worker_routes.py new file mode 100644 index 000000000..388f6e8f3 --- /dev/null +++ b/tests/test_kubeflow_worker_routes.py @@ -0,0 +1,230 @@ +"""Tests for task-bound Kubeflow Worker protocol behavior.""" + +from __future__ import annotations + +import unittest +from unittest.mock import AsyncMock, MagicMock, patch + +from fastapi import HTTPException + +from simpletuner.simpletuner_sdk.server.models.worker import Worker, WorkerStatus, WorkerType +from simpletuner.simpletuner_sdk.server.routes.workers import ( + HeartbeatRequest, + JobStatusUpdate, + WorkerRegistrationRequest, + register_worker, + update_job_status, + worker_heartbeat, + worker_stream, + worker_streams, +) + + +def _bound_worker() -> Worker: + """Create a provisioned Worker fixture. + + Returns: + Kubeflow Worker bound to one job. + """ + return Worker( + worker_id="worker-1", + name="worker-1", + worker_type=WorkerType.EPHEMERAL, + status=WorkerStatus.CONNECTING, + token_hash="hash", + user_id=1, + provider="kubeflow", + current_job_id="job-1", + ) + + +class KubeflowWorkerRoutesTestCase(unittest.IsolatedAsyncioTestCase): + """Test bound Worker registration, heartbeat, and dispatch.""" + + async def asyncTearDown(self) -> None: + """Remove in-memory Worker streams created by tests.""" + worker_streams.clear() + + async def test_registration_keeps_bound_worker_out_of_idle_pool(self) -> None: + """Verify a provisioned Worker remains assigned while registering.""" + worker = _bound_worker() + worker_repository = AsyncMock() + job_repository = AsyncMock() + job_repository.get.return_value = MagicMock(status="queued") + request = MagicMock() + request.client.host = "10.0.0.2" + + with ( + patch( + "simpletuner.simpletuner_sdk.server.routes.workers.validate_worker_token", + new=AsyncMock(return_value=worker), + ), + patch( + "simpletuner.simpletuner_sdk.server.services.worker_repository.get_worker_repository", + return_value=worker_repository, + ), + patch( + "simpletuner.simpletuner_sdk.server.services.cloud.storage.job_repository.get_job_repository", + return_value=job_repository, + ), + ): + response = await register_worker( + WorkerRegistrationRequest( + name="worker-1", + persistent=False, + provider="kubeflow", + current_job_id="job-1", + ), + request, + "token", + ) + + updates = worker_repository.update_worker.await_args.args[1] + self.assertEqual(updates["status"], WorkerStatus.CONNECTING) + self.assertIsNone(response.abandon_job) + + async def test_stream_dispatches_prebound_job(self) -> None: + """Verify SSE connection triggers targeted dispatch.""" + worker = _bound_worker() + manager = AsyncMock() + + with ( + patch( + "simpletuner.simpletuner_sdk.server.routes.workers.validate_worker_token", + new=AsyncMock(return_value=worker), + ), + patch( + "simpletuner.simpletuner_sdk.server.services.worker_manager.get_worker_manager", + return_value=manager, + ), + ): + await worker_stream("worker-1", "token") + + manager.dispatch_bound_job.assert_awaited_once_with("worker-1") + + async def test_idle_heartbeat_cannot_unbind_provisioned_worker(self) -> None: + """Verify a preassigned Worker cannot advertise itself as idle.""" + worker = _bound_worker() + worker_repository = AsyncMock() + + with ( + patch( + "simpletuner.simpletuner_sdk.server.routes.workers.validate_worker_token", + new=AsyncMock(return_value=worker), + ), + patch( + "simpletuner.simpletuner_sdk.server.services.worker_repository.get_worker_repository", + return_value=worker_repository, + ), + ): + await worker_heartbeat( + HeartbeatRequest( + worker_id="worker-1", + status="idle", + current_job_id=None, + ), + "token", + ) + + updates = worker_repository.update_worker.await_args.args[1] + self.assertEqual(updates["status"], WorkerStatus.CONNECTING) + self.assertEqual(updates["current_job_id"], "job-1") + + + async def test_completed_requires_central_lora_artifact(self) -> None: + """Verify a bound Worker cannot complete before central upload.""" + worker = _bound_worker() + job_repository = AsyncMock() + job_repository.get.return_value = MagicMock( + metadata={"worker_id": worker.worker_id}, + ) + worker_repository = AsyncMock() + service = MagicMock() + service.finalize = AsyncMock() + service.artifacts_received.return_value = False + + with ( + patch( + "simpletuner.simpletuner_sdk.server.routes.workers.validate_worker_token", + new=AsyncMock(return_value=worker), + ), + patch( + "simpletuner.simpletuner_sdk.server.services.cloud.storage.job_repository.get_job_repository", + return_value=job_repository, + ), + patch( + "simpletuner.simpletuner_sdk.server.services.worker_repository.get_worker_repository", + return_value=worker_repository, + ), + patch( + "simpletuner.simpletuner_sdk.server.services.kubeflow_job_service.get_kubeflow_job_service", + return_value=service, + ), + ): + with self.assertRaises(HTTPException) as context: + await update_job_status( + "job-1", + JobStatusUpdate(status="completed"), + "token", + ) + + self.assertEqual(context.exception.status_code, 409) + job_repository.update.assert_not_awaited() + service.finalize.assert_not_awaited() + + async def test_completed_finalizes_after_central_lora_artifact(self) -> None: + """Verify central LoRA receipt unlocks completion and cleanup.""" + worker = _bound_worker() + job_repository = AsyncMock() + job_repository.get.return_value = MagicMock( + metadata={ + "worker_id": worker.worker_id, + "artifact_upload": { + "status": "receiving", + "received_files": [ + "outputs/job-1/tiny-output/pytorch_lora_weights.safetensors" + ], + }, + }, + ) + worker_repository = AsyncMock() + service = MagicMock() + service.finalize = AsyncMock() + service.artifacts_received.return_value = True + + with ( + patch( + "simpletuner.simpletuner_sdk.server.routes.workers.validate_worker_token", + new=AsyncMock(return_value=worker), + ), + patch( + "simpletuner.simpletuner_sdk.server.services.cloud.storage.job_repository.get_job_repository", + return_value=job_repository, + ), + patch( + "simpletuner.simpletuner_sdk.server.services.worker_repository.get_worker_repository", + return_value=worker_repository, + ), + patch( + "simpletuner.simpletuner_sdk.server.services.kubeflow_job_service.get_kubeflow_job_service", + return_value=service, + ), + ): + response = await update_job_status( + "job-1", + JobStatusUpdate(status="completed"), + "token", + ) + + self.assertEqual(response, {"success": True}) + updates = job_repository.update.await_args.args[1] + self.assertEqual(updates["status"], "completed") + self.assertEqual( + updates["metadata"]["artifact_upload"]["status"], + "complete", + ) + service.finalize.assert_awaited_once_with("job-1") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_publishing_providers.py b/tests/test_publishing_providers.py index f291b6165..afb24f9d3 100644 --- a/tests/test_publishing_providers.py +++ b/tests/test_publishing_providers.py @@ -6,6 +6,8 @@ from pathlib import Path from unittest.mock import MagicMock, patch +from botocore.compat import HTTPHeaders + from simpletuner.helpers.publishing.providers.azure_blob import AzureBlobPublishingProvider from simpletuner.helpers.publishing.providers.dropbox import DropboxPublishingProvider from simpletuner.helpers.publishing.providers.s3 import S3PublishingProvider @@ -118,6 +120,33 @@ def fake_download_file(_bucket, _key, filename): self.assertTrue((Path(tmpdir) / "model.safetensors").exists()) self.assertTrue((Path(tmpdir) / "checkpoint_manifest.json").exists()) + def test_server_upload_uses_custom_header_and_single_put(self): + """Server publishing should use bearer headers without multipart calls.""" + config = self.config.copy() + config.update( + { + "request_headers": {"X-SimpleTuner-Secret": "job-token"}, + "force_single_part": True, + } + ) + mock_boto3 = MagicMock() + mock_client = mock_boto3.session.Session.return_value.client.return_value + + with patch.dict(sys.modules, {"boto3": mock_boto3}): + provider = S3PublishingProvider(config) + request = MagicMock(headers=HTTPHeaders()) + header_handler = mock_client.meta.events.register.call_args.args[1] + header_handler(request) + self.assertEqual(request.headers["X-SimpleTuner-Secret"], "job-token") + + with tempfile.TemporaryDirectory() as tmpdir: + artifact = Path(tmpdir) / "model.safetensors" + artifact.write_bytes(b"weights") + provider.publish(artifact) + + mock_client.put_object.assert_called_once() + mock_client.upload_file.assert_not_called() + class TestAzureProvider(unittest.TestCase): def setUp(self): diff --git a/tests/test_s3_put_async.py b/tests/test_s3_put_async.py index 39a5bd993..5d9998b8e 100644 --- a/tests/test_s3_put_async.py +++ b/tests/test_s3_put_async.py @@ -129,12 +129,12 @@ async def test_put_object_with_valid_token(self): ) self.assertEqual(response.status_code, status.HTTP_200_OK) - data = response.json() - - # Check response structure - self.assertIn("ETag", data) - self.assertEqual(data["Key"], "test-file.txt") - self.assertEqual(data["Bucket"], "test-bucket") + self.assertEqual( + response.json(), + {"ETag": response.headers["etag"], "Key": "test-file.txt", "Bucket": "test-bucket"}, + ) + self.assertIn("etag", response.headers) + self.assertTrue(response.headers["etag"].startswith('"')) # Verify file was written upload_dir = Path(self._upload_dir) @@ -314,8 +314,8 @@ async def test_put_object_etag_generation(self): ) # ETags should be the same for identical content - etag1 = response1.json()["ETag"] - etag2 = response2.json()["ETag"] + etag1 = response1.headers["etag"] + etag2 = response2.headers["etag"] self.assertEqual(etag1, etag2) # Different content should have different ETag @@ -325,7 +325,7 @@ async def test_put_object_etag_generation(self): headers={"X-Upload-Token": upload_token}, ) - etag3 = response3.json()["ETag"] + etag3 = response3.headers["etag"] self.assertNotEqual(etag1, etag3) diff --git a/tests/test_s3_upload.py b/tests/test_s3_upload.py index c8e6e2e79..2af47b0a3 100644 --- a/tests/test_s3_upload.py +++ b/tests/test_s3_upload.py @@ -8,6 +8,7 @@ - Object listing """ +import json import tempfile import unittest from pathlib import Path @@ -35,8 +36,10 @@ async def body(self) -> bytes: class MockJob: """Mock job record.""" - def __init__(self, job_id: str = "test-job-123"): + def __init__(self, job_id: str = "test-job-123", provider: str = "test-provider"): self.job_id = job_id + self.provider = provider + self.metadata = {} class TestS3PutObjectAuthentication(unittest.TestCase): @@ -126,9 +129,12 @@ def test_put_with_valid_token_succeeds(self): ): result = asyncio.run(s3_put_object("test-bucket", "test.txt", request)) - self.assertIn("ETag", result) - self.assertEqual(result["Key"], "test.txt") - self.assertEqual(result["Bucket"], "test-bucket") + self.assertEqual(result.status_code, 200) + self.assertIn("etag", result.headers) + self.assertEqual( + json.loads(result.body), + {"ETag": result.headers["etag"], "Key": "test.txt", "Bucket": "test-bucket"}, + ) # Verify file was created file_path = Path(self.temp_dir) / "test-bucket" / "test.txt" @@ -155,7 +161,45 @@ def test_x_simpletuner_secret_header_also_works(self): ): result = asyncio.run(s3_put_object("bucket2", "secret.bin", request)) - self.assertIn("ETag", result) + self.assertEqual(result.status_code, 200) + self.assertIn("etag", result.headers) + self.assertEqual( + json.loads(result.body), + {"ETag": result.headers["etag"], "Key": "secret.bin", "Bucket": "bucket2"}, + ) + + + def test_kubeflow_upload_is_registered_on_the_job(self): + """Test Kubeflow uploads are recorded before job completion.""" + import asyncio + + from simpletuner.simpletuner_sdk.server.routes.cloud.storage import put_object as s3_put_object + + key = "job-456/tiny-output/pytorch_lora_weights.safetensors" + request = MockRequest( + headers={"X-Upload-Token": "valid-token-123"}, + body=b"lora weights", + ) + mock_store = MagicMock() + mock_store.get_job_by_upload_token = AsyncMock( + return_value=MockJob("job-456", provider="kubeflow") + ) + mock_store.update_job = AsyncMock() + + with patch( + "simpletuner.simpletuner_sdk.server.routes.cloud.storage.get_job_store", + return_value=mock_store, + ): + asyncio.run(s3_put_object("outputs", key, request)) + + mock_store.update_job.assert_awaited_once() + job_id, updates = mock_store.update_job.await_args.args + self.assertEqual(job_id, "job-456") + self.assertEqual( + updates["metadata"]["artifact_upload"]["received_files"], + [f"outputs/{key}"], + ) + self.assertEqual(updates["output_url"], "/api/cloud/storage/outputs/job-456") class TestS3PathTraversalPrevention(unittest.TestCase): diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 5ffc21d85..56cb6f27d 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -1902,6 +1902,29 @@ def test_init_resume_checkpoint( mock_logger.info.assert_called() trainer.accelerator.load_state.assert_called_with("/path/to/output/checkpoint-200") + def test_publish_final_artifacts_uploads_output_directory(self): + """Final publishing should run after model files are materialized.""" + trainer = object.__new__(Trainer) + trainer.publishing_manager = Mock(configured=True) + trainer.publishing_manager.providers = [] + trainer.publishing_manager.publish.return_value = [Mock()] + trainer.job_id = "job-123" + trainer.state = {"global_step": 1, "current_epoch": 1} + + with tempfile.TemporaryDirectory() as tmpdir: + output_dir = Path(tmpdir) / "output" + output_dir.mkdir() + (output_dir / "pytorch_lora_weights.safetensors").write_bytes(b"weights") + trainer.config = SimpleNamespace(output_dir=str(output_dir)) + + results = trainer._publish_final_artifacts() + + self.assertEqual(len(results), 1) + args, kwargs = trainer.publishing_manager.publish.call_args + self.assertEqual(args[0], output_dir) + self.assertEqual(kwargs["artifact_name"], "output") + self.assertEqual(kwargs["metadata"]["job_id"], "job-123") + @patch("simpletuner.helpers.training.trainer.PublishingManager") @patch("simpletuner.helpers.training.trainer.AttentionBackendController.on_load_checkpoint") def test_init_resume_checkpoint_downloads_remote(self, mock_attention_backend, mock_manager): diff --git a/tests/test_worker_agent.py b/tests/test_worker_agent.py index 82363c457..1207ed81e 100644 --- a/tests/test_worker_agent.py +++ b/tests/test_worker_agent.py @@ -16,6 +16,11 @@ class WorkerAgentJobStartTestCase(unittest.IsolatedAsyncioTestCase): """Test worker job configuration handoff to the training process.""" + def test_worker_config_has_no_kubeflow_runtime_fields(self) -> None: + """Verify Kubernetes integration does not change Worker configuration.""" + self.assertNotIn("provider", WorkerConfig.__dataclass_fields__) + self.assertNotIn("bound_job_id", WorkerConfig.__dataclass_fields__) + async def test_start_job_uses_json_config_path(self) -> None: """Verify the training process receives the dispatched JSON config.""" job_id = f"test-{uuid.uuid4().hex}" @@ -27,9 +32,11 @@ async def test_start_job_uses_json_config_path(self) -> None: event = { "type": "job_submit", "job_id": job_id, + "provider": "kubeflow", "config": config, "dataloader": None, "upload_endpoint": "/api/cloud/storage", + "upload_token": "test-upload-token", } agent = WorkerAgent( WorkerConfig( @@ -40,7 +47,7 @@ async def test_start_job_uses_json_config_path(self) -> None: ) ) job_dir = Path(f"/tmp/simpletuner_job_{job_id}") - process = Mock() + process = AsyncMock() def close_background_coroutine(coroutine): """Close the mocked monitor coroutine to avoid leaking it. @@ -57,7 +64,11 @@ def close_background_coroutine(coroutine): try: with ( patch.object(agent, "_report_job_status", new=AsyncMock()) as report_status, - patch("simpletuner.worker_agent.subprocess.Popen", return_value=process) as popen, + patch( + "simpletuner.worker_agent.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=process, + ) as create_process, patch("simpletuner.worker_agent.asyncio.create_task", side_effect=close_background_coroutine), ): await agent._start_job(event) @@ -65,15 +76,170 @@ def close_background_coroutine(coroutine): config_path = job_dir / "config.json" self.assertEqual(json.loads(config_path.read_text(encoding="utf-8")), config) - command = popen.call_args.args[0] - process_env = popen.call_args.kwargs["env"] - self.assertEqual(command, [sys.executable, "-m", "simpletuner.train"]) + command = create_process.await_args.args + process_env = create_process.await_args.kwargs["env"] + self.assertEqual(command, (sys.executable, "-m", "simpletuner.train")) self.assertEqual(process_env["CONFIG_PATH"], str(config_path)) self.assertEqual(process_env["SIMPLETUNER_CONFIG_BACKEND"], "json") + self.assertEqual(process_env["SIMPLETUNER_UPLOAD_TOKEN"], "test-upload-token") report_status.assert_awaited_once_with("starting") finally: shutil.rmtree(job_dir, ignore_errors=True) + async def test_start_job_materializes_dataloader_json(self) -> None: + """Verify dispatched dataset content becomes the active config path.""" + job_id = f"test-{uuid.uuid4().hex}" + event = { + "type": "job_submit", + "job_id": job_id, + "provider": "kubeflow", + "config": {"model_family": "sdxl"}, + "dataloader": [{"id": "dataset", "type": "local"}], + "upload_endpoint": "/api/cloud/storage", + } + agent = WorkerAgent( + WorkerConfig( + orchestrator_url="https://orchestrator.example.com", + worker_token="test-token", + name="test-worker", + persistent=True, + ) + ) + job_dir = Path(f"/tmp/simpletuner_job_{job_id}") + + def close_background_coroutine(coroutine): + """Close a mocked monitor coroutine. + + Args: + coroutine: Monitor coroutine created by the Worker. + + Returns: + Mock task placeholder. + """ + coroutine.close() + return Mock() + + try: + with ( + patch.object(agent, "_report_job_status", new=AsyncMock()), + patch( + "simpletuner.worker_agent.asyncio.create_subprocess_exec", + new_callable=AsyncMock, + return_value=AsyncMock(), + ), + patch("simpletuner.worker_agent.asyncio.create_task", side_effect=close_background_coroutine), + ): + await agent._start_job(event) + + dataloader_path = job_dir / "dataloader.json" + written_config = json.loads((job_dir / "config.json").read_text(encoding="utf-8")) + self.assertEqual(json.loads(dataloader_path.read_text(encoding="utf-8")), event["dataloader"]) + self.assertEqual(written_config["data_backend_config"], str(dataloader_path)) + finally: + shutil.rmtree(job_dir, ignore_errors=True) + + async def test_start_legacy_job_preserves_yaml_dataloader(self) -> None: + """Verify ordinary Workers retain their existing YAML launch contract.""" + job_id = f"test-{uuid.uuid4().hex}" + event = { + "type": "job_submit", + "job_id": job_id, + "config": {"model_family": "sdxl"}, + "dataloader": [{"id": "dataset", "type": "local"}], + "upload_endpoint": "/api/cloud/storage", + } + agent = WorkerAgent( + WorkerConfig( + orchestrator_url="https://orchestrator.example.com", + worker_token="test-token", + name="test-worker", + persistent=True, + ) + ) + job_dir = Path(f"/tmp/simpletuner_job_{job_id}") + + def close_background_coroutine(coroutine): + """Close the mocked monitor coroutine to avoid leaking it. + + Args: + coroutine: Monitor coroutine created by the Worker. + + Returns: + A mock task placeholder. + """ + coroutine.close() + return Mock() + + try: + with ( + patch.object(agent, "_report_job_status", new=AsyncMock()), + patch("simpletuner.worker_agent.subprocess.Popen", return_value=Mock()) as popen, + patch("simpletuner.worker_agent.asyncio.create_task", side_effect=close_background_coroutine), + ): + await agent._start_job(event) + + self.assertTrue((job_dir / "dataloader.yaml").exists()) + self.assertFalse((job_dir / "dataloader.json").exists()) + self.assertEqual( + json.loads((job_dir / "config.json").read_text(encoding="utf-8")), + event["config"], + ) + self.assertEqual(popen.call_args.args[0], [sys.executable, "-m", "simpletuner.train"]) + finally: + shutil.rmtree(job_dir, ignore_errors=True) + + async def test_ephemeral_worker_exits_after_training_when_queue_is_empty(self) -> None: + """Verify an ordinary ephemeral Worker keeps the legacy process path.""" + agent = WorkerAgent( + WorkerConfig( + orchestrator_url="https://orchestrator.example.com", + worker_token="test-token", + name="test-worker", + ) + ) + process = Mock() + process.stdout = None + process.poll.return_value = 0 + process.returncode = 0 + agent.training_process = process + agent.current_job = {"job_id": "job-1"} + + with ( + patch.object(agent, "_report_job_status", new=AsyncMock()) as report_status, + patch.object(agent, "_check_queue", new=AsyncMock(return_value=False)) as check_queue, + ): + await agent._monitor_training(Path("/tmp")) + + self.assertTrue(agent.shutdown_requested) + check_queue.assert_awaited_once_with() + self.assertEqual(report_status.await_args_list[0].args, ("training",)) + self.assertEqual(report_status.await_args_list[1].args, ("completed",)) + + async def test_stop_current_job_awaits_async_subprocess(self) -> None: + """Verify Kubeflow cancellation awaits the asyncio subprocess API.""" + agent = WorkerAgent( + WorkerConfig( + orchestrator_url="https://orchestrator.example.com", + worker_token="test-token", + name="test-worker", + ) + ) + process = Mock() + process.returncode = None + process.wait = AsyncMock(return_value=0) + agent.training_process = process + agent.current_job = {"job_id": "job-1"} + + with patch.object(agent, "_report_job_status", new=AsyncMock()) as report_status: + with patch.object(agent, "_is_async_process", return_value=True): + await agent._stop_current_job() + + process.terminate.assert_called_once_with() + process.wait.assert_awaited_once_with() + report_status.assert_awaited_once_with("cancelled") + self.assertIsNone(agent.training_process) + self.assertIsNone(agent.current_job) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_worker_manager.py b/tests/test_worker_manager.py index 2afe1e394..55c0464cf 100644 --- a/tests/test_worker_manager.py +++ b/tests/test_worker_manager.py @@ -36,6 +36,7 @@ def _create_worker( worker_id: str = "worker-1", status: WorkerStatus = WorkerStatus.IDLE, worker_type: WorkerType = WorkerType.PERSISTENT, + provider: str = None, current_job_id: str = None, last_heartbeat: datetime = None, created_at: datetime = None, @@ -46,6 +47,7 @@ def _create_worker( worker_id: Worker ID status: Worker status worker_type: Worker type + provider: Infrastructure provider current_job_id: Current job ID last_heartbeat: Last heartbeat timestamp created_at: Creation timestamp @@ -64,7 +66,7 @@ def _create_worker( token_hash="test-token-hash", user_id=1, gpu_info={"name": "A100", "vram_gb": 80, "count": 1}, - provider=None, + provider=provider, labels={}, current_job_id=current_job_id, last_heartbeat=last_heartbeat, @@ -654,6 +656,62 @@ async def test_health_check_loop_stop_event(self) -> None: await self.worker_manager.stop() + async def test_bound_kubeflow_worker_skips_generic_health_recovery(self) -> None: + """Verify Kubernetes reconciliation owns provisioned Worker failures.""" + worker = self._create_worker( + status=WorkerStatus.CONNECTING, + worker_type=WorkerType.EPHEMERAL, + provider="kubeflow", + current_job_id="job-1", + created_at=datetime.now(timezone.utc) - timedelta(seconds=600), + ) + + await self.worker_manager._check_worker(worker, datetime.now(timezone.utc)) + + self.assertTrue(worker.is_job_bound) + self.mock_worker_repository.update_worker.assert_not_called() + self.mock_worker_repository.delete_worker.assert_not_called() + self.mock_job_store.update_job.assert_not_called() + + async def test_dispatch_bound_job_never_selects_idle_worker(self) -> None: + """Verify a provisioned Worker receives only its preassigned job.""" + worker = self._create_worker( + status=WorkerStatus.CONNECTING, + worker_type=WorkerType.EPHEMERAL, + provider="kubeflow", + current_job_id="job-1", + ) + job = self._create_job( + job_id="job-1", + status="queued", + provider="kubeflow", + metadata={ + "config": {"model_family": "sdxl"}, + "worker_id": "worker-1", + }, + ) + self.mock_worker_repository.get_worker.return_value = worker + self.mock_job_store.get_job.return_value = job + + with ( + patch( + "simpletuner.simpletuner_sdk.server.routes.workers.is_worker_connected", + return_value=True, + ), + patch( + "simpletuner.simpletuner_sdk.server.routes.workers.push_to_worker", + new_callable=AsyncMock, + return_value=True, + ) as push, + ): + dispatched = await self.worker_manager.dispatch_bound_job("worker-1") + + self.assertTrue(dispatched) + self.mock_worker_repository.get_idle_worker_for_job.assert_not_called() + self.assertEqual(push.await_args.args[0], "worker-1") + self.assertEqual(push.await_args.args[1]["job_id"], "job-1") + self.assertEqual(push.await_args.args[1]["provider"], "kubeflow") + if __name__ == "__main__": unittest.main()