diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 7c1549e..739c99d 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -2,16 +2,16 @@ name: Tests on: push: - branches: [ main, v0.2 ] + branches: [ main, v0.2, 0.2.1 ] pull_request: - branches: [ main, v0.2 ] + branches: [ main, v0.2, 0.2.1 ] jobs: test: runs-on: ubuntu-latest strategy: matrix: - python-version: ['3.9', '3.10', '3.11', '3.12', '3.13', '3.14'] + python-version: ['3.10', '3.11', '3.12', '3.13', '3.14'] steps: - name: Checkout code @@ -51,6 +51,10 @@ jobs: run: | python tests/test_security.py + - name: Run script generation security tests + run: | + python tests/test_script_gen_security.py + - name: Run error handling tests run: | python tests/test_error_handling.py diff --git a/.gitignore b/.gitignore index fb0bcf0..fc238f7 100644 --- a/.gitignore +++ b/.gitignore @@ -52,6 +52,13 @@ dmypy.json # pyright .pyright/ +# Testing +.pytest_cache/ +.coverage +.coverage.* +htmlcov/ +.tox/ + # Jupyter Notebook .ipynb_checkpoints @@ -64,9 +71,19 @@ dmypy.json legacy/ *.legacy +# Development features (not ready for release) +coldpress/distributed/ +examples/pytorch_coldpress_run/ + # Generated output directories output/ manifests/ # User-specific files users/test.yaml + +# Documentation +docs/ +#README.md +CHANGELOG.md +CLAUDE.md diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3545ee2..4a9070d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,10 +1,58 @@ # Generated by: Claude Sonnet 4.5 +# Pre-commit hooks for Coldpress +# Install: pip install pre-commit && pre-commit install +# Run manually: pre-commit run --all-files + repos: - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.14.10 hooks: # Run the linter - id: ruff + args: [--fix] + exclude: ^(venv/|legacy/) # Run the formatter - id: ruff-format - args: [--check] + exclude: ^(venv/|legacy/) + + - repo: local + hooks: + # Run validation tests + - id: test-validation + name: validation tests + entry: python tests/test_validation.py + language: system + pass_filenames: false + always_run: true + + # Run security tests + - id: test-security + name: security tests + entry: python tests/test_security.py + language: system + pass_filenames: false + always_run: true + + # Run script generation security tests + - id: test-script-gen-security + name: script generation security tests + entry: python tests/test_script_gen_security.py + language: system + pass_filenames: false + always_run: true + + # Run label tests + - id: test-labels + name: label tests + entry: python tests/test_labels.py + language: system + pass_filenames: false + always_run: true + + # Run error handling tests + - id: test-error-handling + name: error handling tests + entry: python tests/test_error_handling.py + language: system + pass_filenames: false + always_run: true diff --git a/README.md b/README.md index 3fdf651..45c63a6 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,3 @@ - # Coldpress [![Tests](https://github.com/asanaullah/coldpress/workflows/Tests/badge.svg?branch=v0.2)](https://github.com/asanaullah/coldpress/actions/workflows/tests.yml) @@ -8,15 +7,32 @@ Coldpress is a prescriptive manifest generator that reduces the effort and exper **Two-piece architecture:** 1. **Admin** (`coldpress-setup`) - Generates cluster setup manifests (node labels, queues, namespaces, RBAC) -2. **User** (`coldpress`) - Generates job manifests from job specifications, creates JobSet YAML + helper scripts - -## How Does It Work? +2. **User** (`coldpress`) - Transforms vanilla Kubernetes Jobs into orchestrated resources (JobSet, Kubeflow PyTorchJob, KubeRay RayJob, KServe InferenceService) + +## Table of Contents + +- [How It Works](#how-it-works) +- [Installation](#installation) +- [Quick Start](#quick-start) + - [For Administrators](#for-administrators) + - [For Users](#for-users) +- [Core Concepts](#core-concepts) + - [Intent Files](#intent-files) + - [Macros](#macros) +- [Examples](#examples) +- [Testing](#testing) +- [Repository Structure](#repository-structure) +- [Requirements](#requirements) +- [Environment Variables](#environment-variables) +- [License](#license) + +## How It Works ### Overall Flow ``` ┌─────────────────────────────────────────────────────────────────┐ -│ Prerequisites: Kueue and JobSet operators must be installed │ +│ Prerequisites: Operators must be installed (see Requirements) │ └─────────────────────────────────────────────────────────────────┘ ↓ ┌─────────────────────────────────────────────────────────────────┐ @@ -44,28 +60,31 @@ Coldpress is a prescriptive manifest generator that reduces the effort and exper ┌─────────────────────────────────────────────────────────────────┐ │ Phase 2: User Workflow (Repeatable) │ ├─────────────────────────────────────────────────────────────────┤ -│ 1. coldpress generate --config job.yaml → output/job-name/ │ -│ - jobset.yaml (Kubernetes manifest) │ +│ 1. Create job-spec.yaml (vanilla Kubernetes Jobs) │ +│ Create intent.yaml (specify target: jobset|kubeflow|kuberay) │ +│ │ +│ 2. coldpress generate --intent intent.yaml → output/job-name/ │ +│ - Generated manifest (jobset.yaml, pytorchjob.yaml, etc.) │ │ - run.sh, monitor.sh, logs.sh, explore.sh, cp.sh, cleanup.sh │ │ │ -│ 2. User reviews jobset.yaml │ +│ 3. User reviews generated manifest │ │ │ -│ 3. ./run.sh applies JobSet to cluster │ +│ 4. ./run.sh applies manifest to cluster │ │ │ -│ 4. Kueue schedules job when resources available │ +│ 5. Kueue schedules job when resources available │ │ │ -│ 5. Jobs execute: mkdir → task-0 → task-1 → ... │ +│ 6. Jobs execute: mkdir → task-0 → task-1 → ... │ │ - Init containers capture hardware discovery │ │ - Main containers run workload │ │ - Results saved to PVC in task-specific directories │ │ │ -│ 6. ./logs.sh captures logs to PVC │ +│ 7. ./logs.sh captures logs to PVC │ │ │ -│ 7. ./explore.sh opens interactive shell to browse results │ +│ 8. ./explore.sh opens interactive shell to browse results │ │ │ -│ 8. ./cp.sh copies results from PVC (optional) │ +│ 9. ./cp.sh copies results from PVC (optional) │ │ │ -│ 9. ./cleanup.sh deletes JobSet (preserves results in PVC) │ +│ 10. ./cleanup.sh deletes resources (preserves results in PVC) │ └─────────────────────────────────────────────────────────────────┘ ``` @@ -79,20 +98,18 @@ Coldpress is a prescriptive manifest generator that reduces the effort and exper - Outputs timestamped manifests for GitOps workflows **For Users (`coldpress`):** -- Generates JobSet manifests from simple job specifications +- Transforms vanilla Kubernetes Jobs into multiple backend formats: + - **JobSet** - Multi-task workflows with dependencies + - **Kubeflow PyTorchJob** - Distributed PyTorch training + - **KubeRay RayJob** - Ray-based distributed computing + - **KServe InferenceService** - Model serving infrastructure - Configures task dependencies (endpoint blocking, completion blocking) - Configures node affinity rules -- Configures volume mounts and hardware discovery init containers -- Creates helper scripts for job lifecycle management +- Injects discovery init containers for hardware profiling +- Generates helper scripts for job lifecycle management - Validates YAML schemas before generation -**Prerequisites:** -- Kueue and JobSet operators must be installed on the cluster -- Users must exist in cluster authentication system - -## Getting Started - -### Installation +## Installation Choose the installation method that fits your use case: @@ -102,7 +119,7 @@ Choose the installation method that fits your use case: | **Running jobs + dev** | `./setup-env.sh --pipx-editable` | ❌ No | Users who also contribute | | **Development** | `./setup-env.sh --uv` | ✅ Yes | Contributors, testing changes | -#### Quick Start +### Quick Start ```bash # For end users (global install, no activation needed) @@ -115,7 +132,7 @@ source .venv/bin/activate coldpress --version ``` -#### For End Users (pipx) +### For End Users (pipx) **pipx** installs Coldpress in an isolated environment with global CLI access - no activation needed. @@ -141,7 +158,7 @@ pipx uninstall coldpress # Remove completely pipx list # Show installed packages ``` -#### For Developers (uv) +### For Developers (uv) **uv** provides fast, reproducible virtual environments for development work. @@ -166,120 +183,408 @@ coldpress --help source .venv/bin/activate ``` -**Why uv for development?** -- Fast dependency resolution and installation -- Reproducible builds -- Editable install by default (changes reflect immediately) -- Isolated from system Python +## Quick Start ### For Administrators -If you are setting up Coldpress for the first time on a cluster, follow the **[Admin Quickstart Guide](docs/quickstart_admin.md)** to: +If you are setting up Coldpress for the first time on a cluster: + +1. **Generate cluster configuration:** + ```bash + coldpress-setup generate cluster cluster/ocp-test-nerc-mghpcc.yaml + ``` + This creates: + - `manifests/cluster-*.yaml` - ClusterQueue, ResourceFlavors + - `manifests/label-nodes-*.sh` - Node labeling script + +2. **Label nodes and apply cluster config:** + ```bash + ./manifests/label-nodes-*.sh + oc apply -f manifests/cluster-*.yaml + ``` + +3. **Generate project configuration:** + ```bash + coldpress-setup generate project projects/coldpress-project.yaml + ``` + This creates: + - `manifests/project-*.yaml` - Namespace, LocalQueue, PVCs + +4. **Apply project config:** + ```bash + oc apply -f manifests/project-*.yaml + ``` + +5. **Generate user RBAC:** + ```bash + coldpress-setup generate user users/myuser.yaml + ``` + This creates: + - `manifests/user-*.yaml` - RoleBindings for job submission + +6. **Apply user config:** + ```bash + oc apply -f manifests/user-*.yaml + ``` -1. Generate and apply cluster-wide configuration (ClusterQueue, ResourceFlavors) -2. Generate and apply project configuration (namespaces, storage, queues) -3. Generate and apply user RBAC (permissions for job submission) +### For Users -This is a one-time setup process that configures the cluster infrastructure for all users. +Once the admin has completed the cluster setup: + +1. **Create your workload specification:** + ```bash + cd my-workflow/ + # Create job-spec.yaml with vanilla Kubernetes Jobs + # Create intent.yaml specifying target backend and transformations + ``` + +2. **Generate manifest for your chosen backend:** + ```bash + coldpress generate --intent intent.yaml + ``` + This creates (based on `target` in intent.yaml): + - `output/job-name/jobset.yaml` (target: jobset) + - `output/job-name/pytorchjob.yaml` (target: kubeflow) + - `output/job-name/rayjob.yaml` (target: kuberay) + - `output/job-name/inferenceservice.yaml` (target: kserve) + - Helper scripts: `run.sh`, `monitor.sh`, `logs.sh`, `explore.sh`, `cleanup.sh` + +3. **Review and apply:** + ```bash + cd output/job-name/ + cat *.yaml # Review generated manifest + ./run.sh # Apply to cluster + ``` + +4. **Monitor job progress:** + ```bash + ./monitor.sh + ``` + +5. **Capture logs:** + ```bash + ./logs.sh + ``` + +6. **Explore results:** + ```bash + ./explore.sh # Opens interactive shell in PVC + ``` + +7. **Copy results locally (optional):** + ```bash + ./cp.sh + ``` + +8. **Clean up:** + ```bash + ./cleanup.sh # Deletes JobSet, preserves results in PVC + ``` + +## Core Concepts + +### Intent Files + +The **intent.yaml** file specifies how Coldpress transforms vanilla Kubernetes Jobs into orchestrated JobSet or Kubeflow resources. + +#### Structure -### For Users +```yaml +# Required +project: +output: +target: jobset | kubeflow | kuberay | kserve # Default: jobset -Once the admin has completed the cluster setup, follow the **[User Quickstart Guide](docs/quickstart_user.md)** to: +# Optional - if omitted, no files mounted +files: + - + - + +# Optional - if omitted, no discovery +discovery: + template: + tasks: all | [task1, task2] + +# Required - must have at least one task +tasks: + - name: # Must match a Job name in job-spec.yaml + replicas: + nodes: [] # Optional + args: + : + env: # Optional - additional environment variables + : + depends_on: # Optional (JobSet only) + task: + wait_for: ready | completion +``` -1. Generate job manifests from your workload specification -2. Review and apply the JobSet to the cluster -3. Monitor job progress and capture logs -4. Explore results in persistent storage -5. Clean up cluster resources (preserves results) +#### Fields + +**Top-Level:** +- `project` (required): Namespace for deployment +- `output` (required): Output directory name +- `target` (optional): Backend to generate (default: `jobset`) + - `jobset` - Multi-task workflows with dependencies + - `kubeflow` - PyTorchJob for distributed PyTorch training + - `kuberay` - RayJob for Ray-based distributed computing + - `kserve` - InferenceService for model serving +- `files` (optional): List of files to mount as ConfigMap +- `discovery` (optional): Discovery configuration +- `tasks` (required): List of task definitions + +**Task Fields:** +- `name` (required): Must match a Job name in job-spec.yaml exactly +- `replicas` (optional): Number of replicas (default: 1) +- `nodes` (optional): Node IDs for pinning +- `args` (optional): Key-value pairs for argument replacement +- `env` (optional): Additional environment variables to inject +- `depends_on` (optional, JobSet only): Dependency specification + +**Dependency Fields:** +- `task` (required): Name of task to depend on +- `wait_for` (required): `ready` (wait for readinessProbe) or `completion` (wait for Job completion) + +#### Example: Multi-Task with Dependencies -This workflow is repeatable for each job you want to run. +```yaml +project: coldpress-project +output: vllm-benchmark +target: jobset + +tasks: + - name: inference-server + replicas: 1 + + - name: benchmark-client + replicas: 1 + depends_on: + task: inference-server + wait_for: ready + args: + target: "http://${inference-server}:8000" +``` -## Documentation +### Macros -- **[Resource Labels](docs/LABELS.md)** - Query and manage Coldpress resources using standard Kubernetes labels - - Find all Coldpress-managed resources: `kubectl get all -A -l app.kubernetes.io/managed-by=coldpress` - - Delete resources by job: `kubectl delete all -n namespace -l coldpress.io/job-id=job-name` -- **[Error Handling](docs/ERROR_HANDLING.md)** - Exit codes, exception handling, and debugging - - Exit code 0: Success, 1: Application errors, 2: Usage errors - - Specific exception types with clear error messages - - Pydantic validation catches configuration errors early -- **Security** - Input validation and injection prevention - - Kubernetes naming rules enforced (lowercase alphanumeric, dashes, dots) - - JSON constructed safely with `json.dumps()` (no f-string injection) - - No hardcoded temp files (timestamped outputs instead) +Macros are placeholders in your job-spec.yaml that Coldpress automatically fills in when generating manifests. -## Quickstart Guides +#### Available Macros -- **[Admin Quickstart](docs/quickstart_admin.md)** - Cluster setup for administrators (one-time) -- **[User Quickstart](docs/quickstart_user.md)** - Running workloads for users (repeatable) +**Task-Local Macros** (current task): +| Macro | Description | Example Value | +|-------|-------------|---------------| +| `${INDEX}` | Replica index within current task | `0`, `1`, `2` | +| `${REPLICAS}` | Total replicas in current task | `2` | +| `${TASK_NAME}` | Name of current task | `ddp-training` | +| `${NODE_ID}` | Physical node ID (if specified) | `1` | +| `${REPLICA_0}` | Pod DNS of replica 0 (current task) | `coldpress-...-task-0-0-0....svc.cluster.local` | +| `${REPLICA_1}` | Pod DNS of replica 1 (current task) | `coldpress-...-task-1-0-0....svc.cluster.local` | -## Repository Structure +**Cross-Task Macros** (reference other tasks): +| Macro | Description | Example | +|-------|-------------|---------| +| `${REPLICA__0}` | Pod DNS of replica 0 of named task | `${REPLICA_inference-server_0}` | +| `${REPLICA__1}` | Pod DNS of replica 1 of named task | `${REPLICA_ddp-training_1}` | +| `${SERVICE_}` | Service DNS for named task (if task has ports) | `${SERVICE_inference-server}` | + +#### Usage Examples + +**Single-Task DDP Training:** +```yaml +tasks: + - name: ddp-training + replicas: 2 + env: + NNODES: "${REPLICAS}" + MASTER_ADDR: "${REPLICA_ddp-training_0}" + RANK: "${INDEX}" ``` -coldpress/ -├── coldpress/ # CLI: Job manifest generator -├── coldpress_setup/ # CLI: Cluster setup and configuration -├── coldpress_common/ # Shared validation models (Pydantic) -├── tests/ # Comprehensive test suite -│ ├── test_validation.py # Pydantic model validation tests -│ ├── test_labels.py # Resource labeling tests -│ ├── test_security.py # Security and input validation tests -│ ├── test_error_handling.py # Error handling tests -│ ├── test_roce_disabled.py # RoCE NIC disabled tests -│ ├── test_exit_codes.sh # Shell exit code tests -│ └── run_all_tests.sh # Run full test suite -├── discovery/ # Hardware discovery pod templates -├── projects/ # Example project configs (namespace, storage) -├── examples/ # Example workloads (config.yaml + job-spec.yaml) -├── cluster/ # Example cluster-wide configurations -├── users/ # Example user RBAC configurations -├── docs/ # Documentation -├── pyproject.toml # Package configuration (modern Python packaging) -├── setup.py # Package setup (legacy, for backward compatibility) -└── setup-env.sh # Environment setup script + +**Multi-Task Client-Server:** +```yaml +# job-spec.yaml +--- +apiVersion: batch/v1 +kind: Job +metadata: + name: inference-server +spec: + template: + spec: + containers: + - name: server + ports: + - containerPort: 8000 + env: + - name: VLLM_PORT + value: "8000" + +--- +apiVersion: batch/v1 +kind: Job +metadata: + name: benchmark-client +spec: + template: + spec: + containers: + - name: client + env: + - name: GUIDELLM_TARGET + value: "${SERVICE_inference-server}" # Resolved automatically ``` -## Testing +## Examples + +Coldpress supports multiple backend targets. Each example demonstrates a different use case and backend. -Coldpress includes a comprehensive test suite that validates all fixes for [GitHub issue #37](https://github.com/asanaullah/coldpress/issues/37). +### PyTorch DDP Training (JobSet) -### Run All Tests +Distributed PyTorch training with 2 workers across 2 GPUs using JobSet. +**Run:** ```bash -# From repository root -./tests/run_all_tests.sh +coldpress generate --intent examples/pytorch_ddp_training/intent_jobset.yaml +cd output/ddp-training-job/ +./run.sh ``` -### Run Individual Tests +**What it demonstrates:** +- Multi-replica DDP training with automatic DNS coordination +- ConfigMap mounting for training script +- Hardware discovery via init containers +- Persistent storage for checkpoints + +### PyTorch DDP Training (Kubeflow) +Same training workload using Kubeflow's PyTorchJob operator. + +**Run:** ```bash -python tests/test_validation.py # Pydantic model validation -python tests/test_labels.py # Resource labels -python tests/test_security.py # Security & input validation -python tests/test_error_handling.py # Error handling & exit codes -python tests/test_roce_disabled.py # RoCE NIC disabled -bash tests/test_exit_codes.sh # Shell exit codes +coldpress generate --intent examples/pytorch_ddp_training/intent_kubeflow.yaml +cd output/ddp-training-job/ +./run.sh ``` -### CI/CD +**What it demonstrates:** +- PyTorchJob CRD for native PyTorch distributed training +- Automatic MASTER_ADDR, MASTER_PORT, RANK injection by Kubeflow +- Single vanilla job-spec.yaml works across both JobSet and Kubeflow targets + +### Ray Distributed Training (KubeRay) + +Ray-based distributed training using KubeRay operator. + +**Run:** +```bash +coldpress generate --intent examples/pytorch_ray_training/intent_kuberay.yaml +cd output/ray-training-job/ +./run.sh +``` + +**What it demonstrates:** +- RayJob CRD for Ray-based workloads +- Automatic Ray cluster setup (head + worker nodes) +- Resource scaling via replicas (2 pods → 4 GPUs total) + +### vLLM + GuideLLM Benchmark (JobSet) + +Multi-task client-server workflow with dependency management. + +**Run:** +```bash +coldpress generate --intent examples/vllm_guidellm_benchmark/intent_jobset.yaml +cd output/vllm-benchmark-job/ +./run.sh +``` + +**What it demonstrates:** +- Task dependencies (`wait_for: ready`) +- Service discovery via `${SERVICE_*}` macros +- Automatic service creation for tasks with ports +- Client waits for server readiness before starting + +**Comparison with manual approach:** +- Manual: 100+ lines of bash for orchestration, polling, error handling +- Coldpress: 14 lines of YAML (intent file) -Tests run automatically on every push and pull request via GitHub Actions. The workflow tests on Python 3.9, 3.10, 3.11, 3.12, 3.13, and 3.14. +### vLLM Inference (KServe) -See `tests/README.md` for detailed test documentation. +Model serving using KServe InferenceService. + +**Run:** +```bash +coldpress generate --intent examples/vllm_guidellm_benchmark/intent_kserve.yaml +cd output/vllm-kserve-inference/ +./run.sh +``` + +**What it demonstrates:** +- KServe InferenceService for production model serving +- Single job-spec.yaml reused across JobSet and KServe targets +- Automatic scaling and traffic management via KServe + +## Testing + +Comprehensive test suite validates: +- Pydantic model validation (`test_validation.py`) +- Standard Kubernetes labels (`test_labels.py`) +- Security and input validation (`test_security.py`) +- Error handling (`test_error_handling.py`) +- Exit codes (`test_exit_codes.sh`) + +**Run all tests:** +```bash +./tests/run_all_tests.sh +``` + +**Run individual tests:** +```bash +python tests/test_validation.py +python tests/test_labels.py +bash tests/test_exit_codes.sh +``` + +Tests run automatically via GitHub Actions on every push. + +## Repository Structure + +``` +coldpress/ +├── coldpress/ # CLI: Job manifest generator +├── coldpress_setup/ # CLI: Cluster setup and configuration +├── coldpress_common/ # Shared validation models (Pydantic) +├── tests/ # Comprehensive test suite +├── discovery/ # Hardware discovery pod templates +├── projects/ # Example project configs (namespace, storage) +├── examples/ # Example workloads (intent.yaml + job-spec.yaml) +├── cluster/ # Example cluster-wide configurations +├── users/ # Example user RBAC configurations +├── docs/ # Documentation (CHANGELOG, quickstart guides) +├── pyproject.toml # Package configuration (modern Python packaging) +├── setup.py # Package setup (legacy, for backward compatibility) +└── setup-env.sh # Environment setup script +``` ## Requirements **Cluster:** - Kubernetes cluster (tested on OpenShift 4.21.5, Kubernetes v1.34.4) -- Kueue operator (tested with v0.11.6, API v1beta1) -- JobSet operator (tested with v1.0.0, API v1alpha2) +- Kueue operator (tested with v0.11.6, API v1beta1) - required for all targets +- **Additional operators** (depending on target backend): + - JobSet operator (v1.0.0+, API v1alpha2) - for `target: jobset` + - Kubeflow Training Operator (v1.8+) - for `target: kubeflow` + - KubeRay operator (v1.0+) - for `target: kuberay` + - KServe (v0.11+) - for `target: kserve` **Local development:** -- Python 3.9+ (tested on Python 3.14) +- Python 3.10+ (tested on Python 3.14) **Cluster tools:** -- `kubectl` or `oc` CLI (tested with oc 4.17.0) +- `oc` CLI (tested with oc 4.17.0) ## Environment Variables @@ -298,7 +603,7 @@ Customize directory locations with environment variables: **Example:** ```bash export COLDPRESS_OUTPUT_DIR=jobs -coldpress generate --config examples/pytorch_ddp_training/config.yaml +coldpress generate --intent examples/pytorch_ddp_training/intent_jobset.yaml # Outputs to: jobs/ddp-training-job/ instead of output/ddp-training-job/ export COLDPRESS_MANIFESTS_DIR=gitops/manifests @@ -306,109 +611,24 @@ coldpress-setup generate project coldpress-project.yaml # Outputs to: gitops/manifests/project-coldpress-project-*.yaml ``` -## Example: PyTorch DDP Training - -**Job spec** (examples/pytorch_ddp_training/job-spec.yaml): - -Note: This generates a JobSet named `coldpress-ddp-training` in the cluster. - -```yaml -name: ddp-training -tolerate_all: true - -containers: - - name: training - image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime - workingDir: /workspace - command: ["python", "-m", "torch.distributed.run"] - args: - - --nproc_per_node=2 - - --nnodes=1 - - train.py - - --dataset=mnist - - --train-test-split=0.8 - - --epochs=50 - - --batch-size=128 - - --hidden-size=4096 - - --lr=0.01 - - --output-dir=/results/checkpoints - resources: - requests: - nvidia.com/gpu: "2" - memory: "16Gi" - cpu: "8" - limits: - nvidia.com/gpu: "2" - memory: "16Gi" - env: - - name: NCCL_DEBUG - value: "INFO" - -volumes: - - name: results - mount: /results - - name: dshm - type: emptyDir - medium: Memory - sizeLimit: 16Gi - mount: /dev/shm -``` - -**Config** (examples/pytorch_ddp_training/config.yaml): -```yaml -project: coldpress-project - -# Discovery - runs as init container per task to capture actual node hardware -discovery: user_snapshot # Simple format -# Or use detailed format: -# discovery: -# template: user_snapshot -# tasks: all # or [0, 1] for specific tasks +## Resource Labels -output: ddp-training-job +All generated resources include standard Kubernetes labels for easy querying and management: -# Files to mount into container (creates ConfigMap) -files: - - train.py - - model_config.json -``` +- `app.kubernetes.io/managed-by: coldpress` - Identifies all Coldpress-managed resources +- `app.kubernetes.io/version: 0.2.1` - Tracks Coldpress version +- `coldpress.io/job-id: {job_name}` - Job-specific identifier for compute resources -**Generate and run:** +**Query all Coldpress resources:** ```bash -coldpress generate --config examples/pytorch_ddp_training/config.yaml -cd output/ddp-training-job/ -./run.sh -``` - -**Results structure in PVC:** +oc get all -A -l app.kubernetes.io/managed-by=coldpress ``` -/data/coldpress-project/coldpress_results/ddp-training-{uid}-{timestamp}/ -├── task-0/ -│ ├── discovery_user_snapshot.json # Hardware/benchmark data for task 0 -│ ├── checkpoints/ -│ │ ├── model_weights.pth # Trained model -│ │ └── training_stats.json # Training metrics -└── logs/ - ├── {pod-name}.log # Individual pod logs - └── combined.log # Combined logs -``` - -## Example: vLLM + GuideLLM Benchmark - -Multi-task workflow with endpoint blocking: +**Delete resources by job:** ```bash -coldpress generate --config examples/vllm_guidellm_benchmark/config.yaml -cd output/vllm-benchmark-job/ -./run.sh +oc delete all -n namespace -l coldpress.io/job-id=job-name ``` -The job-spec defines: -- **Task 1**: vLLM inference server with readinessProbe (endpoint blocking) -- **Task 2**: GuideLLM benchmark client that waits for server readiness - -See [examples/README.md](examples/README.md) for more details. - ## License See LICENSE file. diff --git a/cluster/ocp-test-nerc-mghpcc.yaml b/cluster/ocp-test-nerc-mghpcc.yaml index 76e5753..64d3d90 100644 --- a/cluster/ocp-test-nerc-mghpcc.yaml +++ b/cluster/ocp-test-nerc-mghpcc.yaml @@ -7,10 +7,10 @@ # Node IDs are assigned based on position in list (0, 1, 2, ...) nodes: - hostname: wrk-4 # Will be labeled coldpress.node=0 - gpus: 4 + gpus: 2 roce_nics: 4 - hostname: wrk-6 # Will be labeled coldpress.node=1 - gpus: 4 + gpus: 2 roce_nics: 4 # ClusterQueue name @@ -30,6 +30,8 @@ spec: - "BatchJob" - "JobSet" - "Pod" + - "PyTorchJob" + - "RayJob" --- # JobSet Operator apiVersion: operator.openshift.io/v1 @@ -42,7 +44,7 @@ spec: operatorLogLevel: Normal --- # ResourceFlavor for node0 -apiVersion: kueue.x-k8s.io/v1beta1 +apiVersion: kueue.x-k8s.io/v1beta2 kind: ResourceFlavor metadata: name: node0 @@ -51,7 +53,7 @@ spec: coldpress.node: "0" --- # ResourceFlavor for node1 -apiVersion: kueue.x-k8s.io/v1beta1 +apiVersion: kueue.x-k8s.io/v1beta2 kind: ResourceFlavor metadata: name: node1 @@ -60,7 +62,7 @@ spec: coldpress.node: "1" --- # ClusterQueue for Coldpress -apiVersion: kueue.x-k8s.io/v1beta1 +apiVersion: kueue.x-k8s.io/v1beta2 kind: ClusterQueue metadata: name: coldpress-cluster-queue @@ -76,7 +78,7 @@ spec: - name: memory nominalQuota: 1024Gi - name: nvidia.com/gpu - nominalQuota: 4 + nominalQuota: 2 - name: openshift.io/eno5np0rdma nominalQuota: 1 - name: openshift.io/eno6np0rdma @@ -92,7 +94,7 @@ spec: - name: memory nominalQuota: 1024Gi - name: nvidia.com/gpu - nominalQuota: 4 + nominalQuota: 2 - name: openshift.io/eno5np0rdma nominalQuota: 1 - name: openshift.io/eno6np0rdma diff --git a/coldpress/__init__.py b/coldpress/__init__.py index b59d2ff..cd3a266 100644 --- a/coldpress/__init__.py +++ b/coldpress/__init__.py @@ -1,4 +1,4 @@ -# Generated by: Claude Sonnet 4.5 +# Assisted by: Claude Sonnet 4.5 """Coldpress - AI/HPC workload orchestration for Kubernetes/OpenShift.""" -__version__ = "0.2.0" +__version__ = "0.2.1" diff --git a/coldpress/cli.py b/coldpress/cli.py index d790c8e..a47fc26 100644 --- a/coldpress/cli.py +++ b/coldpress/cli.py @@ -1,109 +1,114 @@ -# Generated by: Claude Sonnet 4.5 -"""Coldpress CLI - Generate JobSet manifests and scripts for AI/HPC workloads.""" +# Assisted by: Claude Sonnet 4.5 +"""Coldpress CLI - Generate JobSet and Kubeflow manifests for AI/HPC workloads. + +Configuration Priority: + 1. CLI arguments (--intent, --output-dir, etc.) - highest priority + 2. Environment variables (COLDPRESS_*_DIR) + 3. Default values - lowest priority + +Environment Variables: + COLDPRESS_DISCOVERY_DIR: Directory for cluster discovery configs (default: "discovery") + COLDPRESS_PROJECT_DIR: Directory for project configs (default: "projects") + COLDPRESS_OUTPUT_DIR: Directory for generated output (default: "output") +""" import click import yaml import os import json +import shutil +import traceback from datetime import datetime, timezone -from .generator import generate_jobset, jobset_to_yaml, services_to_yaml +from .jobset_generator import ( + generate_jobset_from_intent, + jobset_to_yaml, + services_to_yaml, +) +from .kubeflow_generator import ( + generate_pytorchjob_from_intent, + generate_inferenceservice_from_intent, + kubeflow_to_yaml, +) +from .kuberay_generator import ( + generate_rayjob_from_intent, + kuberay_to_yaml, +) from .script_gen import write_scripts +from .constants import ( + COLDPRESS_VERSION, + BACKEND_NAME_MAP, + DEFAULT_PROJECT_DIR, + DEFAULT_OUTPUT_DIR, +) +from .utils import extract_configmap_name, get_storage_pvc_name from coldpress_common import ( - validate_config, + validate_intent, validate_project_config, - validate_task_specs, ) from pydantic import ValidationError -# Default directories (can be overridden with environment variables) -DISCOVERY_DIR = os.getenv("COLDPRESS_DISCOVERY_DIR", "discovery") -PROJECT_DIR = os.getenv("COLDPRESS_PROJECT_DIR", "projects") -OUTPUT_DIR = os.getenv("COLDPRESS_OUTPUT_DIR", "output") - @click.group() -@click.version_option(version="0.2.0") +@click.version_option(version=COLDPRESS_VERSION) def cli(): """Coldpress - AI/HPC workload orchestration for Kubernetes/OpenShift.""" pass -def _load_and_validate_config( - config_path, project_override, discovery_override, output_override -): - """Load and validate configuration from file and CLI args.""" - with open(config_path, "r") as f: - config_data = yaml.safe_load(f) +def _load_intent(intent_path): + """Load and validate intent.yaml.""" + if not os.path.exists(intent_path): + raise FileNotFoundError(f"Could not find intent.yaml at {intent_path}") + + with open(intent_path, "r") as f: + intent_data = yaml.safe_load(f) - # Validate config schema + # Validate intent schema try: - validated_config = validate_config(config_data) + validated_intent = validate_intent(intent_data) except ValidationError as e: - raise ValueError(f"Config validation failed: {e}") from e - - config_dir = os.path.dirname(os.path.abspath(config_path)) - project = project_override or validated_config.project - - # Handle discovery - CLI override or config (could be string or DiscoveryConfig) - if discovery_override: - # CLI override is always a simple string, convert to DiscoveryConfig format - discovery = {"template": discovery_override, "tasks": "all"} - elif validated_config.discovery: - # From config file - already validated as DiscoveryConfig - if isinstance(validated_config.discovery, str): - discovery = {"template": validated_config.discovery, "tasks": "all"} - else: - discovery = { - "template": validated_config.discovery.template, - "tasks": validated_config.discovery.tasks, - } - else: - discovery = None - - output_name = output_override or validated_config.output - nodes_from_config = validated_config.nodes + raise ValueError(f"Intent validation failed: {e}") from e - if not project: - raise ValueError("Must provide --project or specify 'project' in config") + return validated_intent - return config_data, config_dir, project, discovery, output_name, nodes_from_config - -def _load_task_specs(config_dir): - """Load and validate task specifications from job-spec.yaml.""" - job_spec_path = os.path.join(config_dir, "job-spec.yaml") +def _load_vk8s_jobs(job_spec_path): + """Load vanilla Kubernetes Jobs from job-spec.yaml.""" if not os.path.exists(job_spec_path): - raise FileNotFoundError(f"Could not find job-spec.yaml in {config_dir}") + raise FileNotFoundError(f"Could not find job-spec.yaml at {job_spec_path}") with open(job_spec_path, "r") as f: - task_specs_data = list(yaml.safe_load_all(f)) + manifests = list(yaml.safe_load_all(f)) - if not task_specs_data: - raise ValueError("Job spec file is empty") + if not manifests: + raise ValueError("job-spec.yaml is empty") - # Validate task specifications - try: - validated_tasks = validate_task_specs(task_specs_data) - except ValidationError as e: - raise ValueError(f"Task spec validation failed: {e}") from e + # Extract Jobs and pass through other resources + jobs = {} + other_resources = [] - # Convert back to dict for compatibility with rest of code - return [ - task.model_dump(by_alias=True, exclude_none=True) for task in validated_tasks - ] + for manifest in manifests: + if manifest.get("kind") == "Job" and manifest.get("apiVersion") == "batch/v1": + job_name = manifest.get("metadata", {}).get("name") + if not job_name: + raise ValueError("Job manifest missing metadata.name") + jobs[job_name] = manifest + else: + # Pass through non-Job resources (Services, ConfigMaps, etc.) + other_resources.append(manifest) + if not jobs: + raise ValueError("No batch/v1 Jobs found in job-spec.yaml") -def _determine_job_name(task_specs): - """Determine job name from task specifications.""" - if len(task_specs) == 1: - return task_specs[0]["name"] - return f"{task_specs[0]['name']}-workflow" + return jobs, other_resources def _load_project_config(project_name): """Load project configuration file.""" - project_config_file = os.path.join(PROJECT_DIR, f"{project_name}.yaml") + # Read config at call time to enable testing and runtime reconfiguration + project_dir = os.getenv("COLDPRESS_PROJECT_DIR", DEFAULT_PROJECT_DIR) + project_config_file = os.path.join(project_dir, f"{project_name}.yaml") if not os.path.exists(project_config_file): raise FileNotFoundError(f"Project config not found: {project_config_file}") @@ -123,157 +128,24 @@ def _load_project_config(project_name): return project_config, namespace -def _prepare_configmap_files(config_dir, files_list, job_name): - """Read and prepare files for ConfigMap.""" - if not files_list: - return None - - configmap_name = f"coldpress-{job_name}-files" - file_data = {} - for file_name in files_list: - file_path = os.path.join(config_dir, file_name) - if not os.path.exists(file_path): - raise FileNotFoundError(f"File not found: {file_path}") - with open(file_path, "r") as f: - file_data[file_name] = f.read() - - return { - "name": configmap_name, - "files": list(file_data.keys()), - "data": file_data, - } - - -def _allocate_nodes_for_tasks(task_specs, manual_nodes): - """Allocate nodes for each task.""" - node_assignments = {} - - if manual_nodes: - # User provided explicit node assignments - for task_id, node_id in enumerate(manual_nodes): - if task_id < len(task_specs): - node_assignments[task_id] = node_id - return node_assignments - - # Default: let Kubernetes scheduler decide (use "any" coldpress node) - for task_id in range(len(task_specs)): - node_assignments[task_id] = "any" - - return node_assignments - - -def _write_output_files( - output_dir, - jobset, - services, - job_spec, - base_dir, - config_dir, - task_specs, - node_assignments, -): - """Write all generated files to output directory.""" - os.makedirs(output_dir, exist_ok=True) - - # Write JobSet YAML - jobset_file = os.path.join(output_dir, "jobset.yaml") - with open(jobset_file, "w") as f: - f.write(jobset_to_yaml(jobset)) - - # Write Services YAML if any - if services: - services_file = os.path.join(output_dir, "services.yaml") - with open(services_file, "w") as f: - f.write(services_to_yaml(services)) - - # Copy ConfigMap files - configmap_name = None - configmap_files = [] - if "configmap" in job_spec: - import shutil - - cm_info = job_spec["configmap"] - configmap_name = cm_info["name"] - for file_name in cm_info["files"]: - src_path = os.path.join(config_dir, file_name) - dst_path = os.path.join(output_dir, file_name) - shutil.copy(src_path, dst_path) - configmap_files.append(file_name) - - # Write metadata - metadata = { - "job_name": job_spec["name"], - "namespace": job_spec["namespace"], - "generated_at": datetime.now(timezone.utc).isoformat(), - "base_dir": base_dir, - "node_assignments": node_assignments, - "tasks": [ - {"name": t.get("name"), "containers": len(t.get("containers", []))} - for t in task_specs - ], - } - metadata_file = os.path.join(output_dir, "metadata.json") - with open(metadata_file, "w") as f: - json.dump(metadata, f, indent=2) - - # Generate bash scripts - storage_pvc = job_spec.get("storage", {}).get( - "results", f"coldpress-{job_spec['namespace']}-storage" - ) - write_scripts( - output_dir, - f"coldpress-{job_spec['name']}", - job_spec["namespace"], - storage_pvc, - base_dir, - configmap_name, - configmap_files, - ) - - return output_dir - - @cli.command() @click.option( - "-c", - "--config", + "-i", + "--intent", required=True, - help="Config file specifying workloads, project, discovery, and output", + help="Intent file specifying transformations (intent.yaml)", type=click.Path(exists=True), ) -@click.option( - "-p", - "--project", - help="Override project from config (loads from $COLDPRESS_PROJECT_DIR/{project}.yaml)", - type=str, -) -@click.option( - "-d", - "--discovery", - help="Override discovery from config (loads from discovery/{name}.yaml)", - type=str, -) -@click.option( - "-o", "--output", help="Override output directory from config", type=click.Path() -) -@click.option( - "--node", - multiple=True, - type=int, - help="Manually assign tasks to specific nodes (--node 0 --node 1 ...)", -) -@click.option( - "--file", - multiple=True, - type=str, - help="Add files to mount into container (--file train.py --file config.json ...)", -) -def generate(config, project, discovery, output, node, file): +def generate(intent): """ - Generate JobSet YAML and bash scripts from config file. + Generate JobSet or Kubeflow YAML manifests from intent.yaml. + + The intent file specifies: + - Target backend (jobset or kubeflow) + - Transformations to apply to vanilla k8s job-spec.yaml + - Task replication and dependencies - The config file specifies project, discovery template, and output directory. - The CLI auto-discovers job-spec.yaml in the same directory as the config file. + The CLI auto-discovers job-spec.yaml in the same directory as intent.yaml. Environment Variables: COLDPRESS_DISCOVERY_DIR - Discovery templates directory (default: discovery) @@ -281,112 +153,169 @@ def generate(config, project, discovery, output, node, file): COLDPRESS_OUTPUT_DIR - Default output directory (default: output) Examples: - coldpress generate --config examples/pytorch_ddp_training/config.yaml - coldpress generate -c examples/vllm_guidellm_benchmark/config.yaml -p different-project - coldpress generate -c my-workflow/config.yaml -o custom-output + coldpress generate --intent examples/pytorch_ddp_training/intent.yaml + coldpress generate -i examples/vllm_guidellm_benchmark/intent.yaml """ try: - config_data, config_dir, project, discovery, output_name, nodes_from_config = ( - _load_and_validate_config(config, project, discovery, output) - ) - task_specs = _load_task_specs(config_dir) - job_name = _determine_job_name(task_specs) - output_path = os.path.join(OUTPUT_DIR, output_name or job_name) + # Read config at call time to enable testing and runtime reconfiguration + output_dir = os.getenv("COLDPRESS_OUTPUT_DIR", DEFAULT_OUTPUT_DIR) - project_config, namespace = _load_project_config(project) + # Load and validate intent.yaml + intent_config = _load_intent(intent) + intent_dir = os.path.dirname(os.path.abspath(intent)) - # Build job spec - job_spec = { - "name": job_name, - "tasks": task_specs, - "namespace": namespace, - "storage": project_config.get("storage", {}), - } + # Load vanilla k8s Jobs from job-spec.yaml + job_spec_path = os.path.join(intent_dir, "job-spec.yaml") + vk8s_jobs, other_resources = _load_vk8s_jobs(job_spec_path) - # Add discovery configuration - if discovery: - template_name = discovery["template"] - discovery_template_path = f"{DISCOVERY_DIR}/{template_name}.yaml" - if not os.path.exists(discovery_template_path): - raise FileNotFoundError( - f"Discovery template not found: {discovery_template_path}" - ) - job_spec["discovery"] = { - "template": discovery_template_path, - "tasks": discovery["tasks"], - } - - # Prepare ConfigMap files - files_list = list(config_data.get("files", [])) + list(file) - configmap_info = _prepare_configmap_files(config_dir, files_list, job_name) - if configmap_info: - job_spec["configmap"] = configmap_info + # Load project config + project_config, namespace = _load_project_config(intent_config.project) + + # Output path + output_path = os.path.join(output_dir, intent_config.output) # Display generation info - click.echo(f"Generating JobSet for: {job_name}") - click.echo(f"Project: {project}") + click.echo("Generating manifests") + click.echo(f"Project: {intent_config.project}") click.echo(f"Namespace: {namespace}") - if discovery: - tasks_info = discovery["tasks"] - if tasks_info == "all": - click.echo(f"Discovery: {discovery['template']} (all tasks)") - else: - click.echo(f"Discovery: {discovery['template']} (tasks: {tasks_info})") - click.echo(f"Tasks: {len(task_specs)}") - - # Allocate nodes (priority: CLI --node > config nodes > default scheduler) - manual_nodes = None - if node: - # CLI --node flag takes precedence - manual_nodes = list(node) - click.echo( - f"Using node assignments from CLI: {dict(enumerate(manual_nodes))}" + click.echo(f"Tasks: {len(intent_config.tasks)}") + for task in intent_config.tasks: + replicas = task.replicas or 1 + nodes_info = f" (nodes: {task.nodes})" if task.nodes else "" + click.echo(f" - {task.name}: {replicas} replica(s){nodes_info}") + + # Use backend from intent.yaml + backend = intent_config.target + backend_name = BACKEND_NAME_MAP.get(backend, "JobSet") + click.echo(f"Target: {backend_name}") + + # Generate manifests based on target + if backend == "kubeflow": + # Generate Kubeflow manifest (PyTorchJob) + manifest, base_dir = generate_pytorchjob_from_intent( + vk8s_jobs, intent_config, project_config, namespace + ) + manifest_type = "pytorchjob" + services = [] # Kubeflow doesn't use separate services + elif backend == "kserve": + # Generate KServe manifest (InferenceService) + manifest, base_dir = generate_inferenceservice_from_intent( + vk8s_jobs, intent_config, project_config, namespace ) - elif nodes_from_config: - # Use nodes from config file - manual_nodes = nodes_from_config - click.echo( - f"Using node assignments from config: {dict(enumerate(manual_nodes))}" + manifest_type = "inferenceservice" + services = [] # KServe doesn't use separate services + elif backend == "kuberay": + # Generate KubeRay manifest (RayJob) + manifest, base_dir = generate_rayjob_from_intent( + vk8s_jobs, intent_config, project_config, namespace ) + manifest_type = "rayjob" + services = [] # KubeRay doesn't use separate services else: - # Default: let Kubernetes scheduler decide - click.echo( - "Node scheduling: Kubernetes will select any coldpress-labeled node" + # Generate JobSet manifest + manifest, services, base_dir = generate_jobset_from_intent( + vk8s_jobs, intent_config, project_config, namespace ) - - node_assignments = _allocate_nodes_for_tasks(task_specs, manual_nodes) - - # Display task info - if not manual_nodes: - for task_id, task in enumerate(task_specs): - req_gpus = sum( - int( - c.get("resources", {}) - .get("requests", {}) - .get("nvidia.com/gpu", "0") - ) - for c in task.get("containers", []) - ) - click.echo( - f" Task {task_id} ({task.get('name', f'task-{task_id}')}) - GPUs: {req_gpus} (scheduler will pick node)" - ) - - # Generate JobSet - jobset, services, base_dir = generate_jobset(job_spec, node_assignments) - - # Write all output files - output_path = _write_output_files( + manifest_type = "jobset" + + # Write manifests + os.makedirs(output_path, exist_ok=True) + + # Verify output directory is writable + if not os.access(output_path, os.W_OK): + raise PermissionError(f"Output directory is not writable: {output_path}") + + if manifest_type == "jobset": + manifest_file = os.path.join(output_path, "jobset.yaml") + with open(manifest_file, "w") as f: + f.write(jobset_to_yaml(manifest)) + elif manifest_type == "pytorchjob": + manifest_file = os.path.join(output_path, "pytorchjob.yaml") + with open(manifest_file, "w") as f: + f.write(kubeflow_to_yaml(manifest)) + elif manifest_type == "inferenceservice": + manifest_file = os.path.join(output_path, "kservejob.yaml") + with open(manifest_file, "w") as f: + f.write(kubeflow_to_yaml(manifest)) + elif manifest_type == "rayjob": + manifest_file = os.path.join(output_path, "rayjob.yaml") + with open(manifest_file, "w") as f: + f.write(kuberay_to_yaml(manifest)) + + if services: + services_file = os.path.join(output_path, "services.yaml") + with open(services_file, "w") as f: + f.write(services_to_yaml(services)) + + # Copy ConfigMap files + if intent_config.files: + for file_name in intent_config.files: + src_path = os.path.join(intent_dir, file_name) + dst_path = os.path.join(output_path, file_name) + if os.path.exists(src_path): + shutil.copy(src_path, dst_path) + + # Write metadata + metadata = { + "job_name": manifest["metadata"]["name"], + "namespace": namespace, + "backend": manifest_type, + "generated_at": datetime.now(timezone.utc).isoformat(), + "base_dir": base_dir, + "tasks": [ + {"name": task.name, "replicas": task.replicas or 1} + for task in intent_config.tasks + ], + } + metadata_file = os.path.join(output_path, "metadata.json") + with open(metadata_file, "w") as f: + json.dump(metadata, f, indent=2) + + # Determine ConfigMap name and files + configmap_name = None + configmap_files = [] + if intent_config.files: + # Extract ConfigMap name from the manifest using shared utility + configmap_name = extract_configmap_name(manifest, manifest_type) + + # If we found a configmap reference, use the files from intent + if configmap_name: + configmap_files = list(intent_config.files) + + # Generate bash scripts + storage_pvc = get_storage_pvc_name(project_config, namespace) + # Strip coldpress- prefix from manifest name for scripts (script_gen will add it back) + job_name_for_scripts = manifest["metadata"]["name"] + if job_name_for_scripts.startswith("coldpress-"): + job_name_for_scripts = job_name_for_scripts[len("coldpress-") :] + + write_scripts( output_path, - jobset, - services, - job_spec, + job_name_for_scripts, + namespace, + storage_pvc, base_dir, - config_dir, - task_specs, - node_assignments, + configmap_name, + configmap_files, + manifest_type, ) - click.echo(f"\nJob manifest generated successfully in: {output_path}/") + click.echo( + f"\n{backend_name} manifest generated successfully in: {output_path}/" + ) + + # Show manifest file location + if manifest_type == "jobset": + click.echo(" Manifest: jobset.yaml") + elif manifest_type == "pytorchjob": + click.echo(" Manifest: pytorchjob.yaml") + elif manifest_type == "tfjob": + click.echo(" Manifest: tfjob.yaml") + elif manifest_type == "mpijob": + click.echo(" Manifest: mpijob.yaml") + elif manifest_type == "inferenceservice": + click.echo(" Manifest: kservejob.yaml") + click.echo("\nTo run the job:") click.echo(f" cd {output_path}") click.echo(" ./run.sh") @@ -402,10 +331,8 @@ def generate(config, project, discovery, output, node, file): except (ValueError, FileNotFoundError) as e: click.echo(f"Error: {e}", err=True) raise SystemExit(1) - except Exception as e: + except (ValidationError, yaml.YAMLError, OSError, PermissionError) as e: click.echo(f"Error generating JobSet: {e}", err=True) - import traceback - traceback.print_exc() raise SystemExit(1) diff --git a/coldpress/constants.py b/coldpress/constants.py new file mode 100644 index 0000000..39a8da1 --- /dev/null +++ b/coldpress/constants.py @@ -0,0 +1,98 @@ +# Assisted by: Claude Sonnet 4.5 +"""Coldpress constants - centralized configuration values.""" + +# Version +COLDPRESS_VERSION = "0.2.1" + +# Resource naming +COLDPRESS_PREFIX = "coldpress" +COLDPRESS_LABEL_MANAGED_BY = "coldpress" + +# Images +MKDIR_IMAGE = "registry.access.redhat.com/ubi9/ubi-minimal:latest" +EXPLORER_IMAGE = "registry.access.redhat.com/ubi9/ubi-minimal:latest" +COPIER_IMAGE = "registry.access.redhat.com/ubi9/ubi:latest" + +# Labels +COLDPRESS_LABELS = { + "app.kubernetes.io/managed-by": COLDPRESS_LABEL_MANAGED_BY, + "app.kubernetes.io/version": COLDPRESS_VERSION, +} + +# Directory defaults +DEFAULT_DISCOVERY_DIR = "discovery" + +# Script generation defaults +DEFAULT_JOB_TIMEOUT = "1h" +DEFAULT_MASTER_PORT = "29500" +DEFAULT_SLEEP_DURATION = "300" # seconds for helper pods +DEFAULT_SLEEP_INFINITY = "infinity" # for explorer pod + + +# Kueue +def get_kueue_queue_label(namespace: str) -> str: + """Get Kueue local queue label for a namespace.""" + return f"{COLDPRESS_PREFIX}-local-queue-{namespace}" + + +def get_jobset_name(job_id: str) -> str: + """Get JobSet name from job ID.""" + return f"{COLDPRESS_PREFIX}-{job_id}" + + +def get_service_name(jobset_name: str, task_id: int) -> str: + """Get Service name for a task.""" + return f"{COLDPRESS_PREFIX}-s-{jobset_name}-{task_id}" + + +def get_pvc_name(namespace: str) -> str: + """Get default PVC name for a namespace.""" + return f"{COLDPRESS_PREFIX}-{namespace}-storage" + + +# Manifest type configuration for script generation +MANIFEST_CONFIG = { + "jobset": { + "file": "jobset.yaml", + "type": "jobset", + "apply_msg": "JobSet", + "has_services": True, + }, + "pytorchjob": { + "file": "pytorchjob.yaml", + "type": "pytorchjob", + "apply_msg": "PyTorchJob", + }, + "tfjob": { + "file": "tfjob.yaml", + "type": "tfjob", + "apply_msg": "TFJob", + }, + "mpijob": { + "file": "mpijob.yaml", + "type": "mpijob", + "apply_msg": "MPIJob", + }, + "inferenceservice": { + "file": "kservejob.yaml", + "type": "inferenceservice", + "apply_msg": "KServe InferenceService", + }, + "rayjob": { + "file": "rayjob.yaml", + "type": "rayjob", + "apply_msg": "rayjob", + }, +} + +# Backend display names +BACKEND_NAME_MAP = { + "kubeflow": "Kubeflow", + "kserve": "KServe", + "kuberay": "KubeRay", + "jobset": "JobSet", +} + +# Default directories (fallback values when env vars not set) +DEFAULT_PROJECT_DIR = "projects" +DEFAULT_OUTPUT_DIR = "output" diff --git a/coldpress/generator.py b/coldpress/generator.py deleted file mode 100644 index 2ea616b..0000000 --- a/coldpress/generator.py +++ /dev/null @@ -1,875 +0,0 @@ -# Generated by: Claude Sonnet 4.5 -"""JobSet YAML generation for Coldpress jobs.""" - -import os -import sys -import yaml -import hashlib -import shlex -from urllib.parse import urlparse -from datetime import datetime, timezone - -# Standard labels for all Coldpress-managed resources -COLDPRESS_LABELS = { - "app.kubernetes.io/managed-by": "coldpress", - "app.kubernetes.io/version": "0.2.0", -} - - -def generate_configmap(name, namespace, file_data): - """Generate ConfigMap for mounting files into containers. - - Args: - name: ConfigMap name - namespace: Kubernetes namespace - file_data: Dict mapping filename to file content - - Returns: - dict: ConfigMap manifest - """ - return { - "apiVersion": "v1", - "kind": "ConfigMap", - "metadata": { - "name": name, - "namespace": namespace, - "labels": COLDPRESS_LABELS.copy(), - }, - "data": file_data, - } - - -def generate_base_dir(namespace, job_name): - """ - Generate base directory path with uid and timestamp. - - Format: {namespace}/coldpress_results/{job_name}_{uid}_{timestamp} - - Args: - namespace: Kubernetes namespace - job_name: Job name - - Returns: - str: Base directory path - """ - # Generate 8-char hex uid from timestamp - now = datetime.now(timezone.utc) - uid = hashlib.md5(f"{job_name}{now.isoformat()}".encode()).hexdigest()[:8] - timestamp = now.strftime("%Y%m%d_%H%M%S") - - return f"{namespace}/coldpress_results/{job_name}-{uid}-{timestamp}" - - -def _infer_blocking_type_and_health_check(task, task_id, job_id, namespace): - """Infer blocking type from readinessProbe if not explicitly set.""" - if "blocking" in task: - return - - containers = task.get("containers", []) - if containers and "readinessProbe" in containers[0]: - task["blocking"] = "endpoint" - probe = containers[0]["readinessProbe"] - if "httpGet" in probe: - http_get = probe["httpGet"] - path = http_get.get("path", "/") - port = http_get.get("port", 80) - scheme = http_get.get("scheme", "HTTP").lower() - service_name = f"coldpress-s-{job_id}-{task_id}.{namespace}.svc" - task["health_check"] = f"{scheme}://{service_name}:{port}{path}" - else: - task["blocking"] = "completion" - - -def _substitute_dns_in_args(tasks, job_id, namespace): - """Substitute DNS placeholders in container args.""" - task_names = {task.get("name"): i for i, task in enumerate(tasks)} - - for task in tasks: - containers = task.get("containers", []) - if not containers or "args" not in containers[0]: - continue - - args = containers[0]["args"] - for i, arg in enumerate(args): - if not isinstance(arg, str): - continue - - for task_name, target_task_id in task_names.items(): - if f"http://{task_name}:" in arg or f"https://{task_name}:" in arg: - service_name = ( - f"coldpress-s-{job_id}-{target_task_id}.{namespace}.svc" - ) - arg = arg.replace(f"://{task_name}:", f"://{service_name}:") - args[i] = arg - - -def _build_init_jobs(base_dir, data_pvc_name, num_tasks): - """Build mkdir initialization job.""" - jobs = [] - - # Add mkdir job - mkdir_job = build_mkdir_job(base_dir, data_pvc_name, "", num_tasks) - jobs.append(mkdir_job) - previous_job_name = mkdir_job["name"] - previous_blocking_type = "completion" - - return jobs, previous_job_name, previous_blocking_type - - -def _build_task_jobs( - tasks, - job_id, - jobset_name, - namespace, - job_spec, - node_assignments, - data_pvc_name, - model_pvc_name, - base_dir, - previous_job_name, - previous_blocking_type, - discovery_template, - discovery_task_indices, -): - """Build replicated jobs for all tasks. - - Args: - job_id: Base job name for DNS (e.g., "ddp-training") - jobset_name: Full JobSet name with prefix for labels (e.g., "coldpress-ddp-training") - """ - replicated_jobs = [] - services = [] - - for task_id, task in enumerate(tasks): - node_id = str(node_assignments.get(task_id, 0)) - - container_spec = build_container_spec( - task, task_id, job_id, namespace, data_pvc_name, model_pvc_name, base_dir - ) - - pod_spec = build_pod_spec( - task, - task_id, - job_id, - jobset_name, - node_id, - container_spec, - data_pvc_name, - job_spec.get("configmap"), - base_dir, - ) - - # Add discovery init container if this task should run discovery - if task_id in discovery_task_indices and discovery_template: - init_container = build_discovery_init_container( - discovery_template, task_id, base_dir, data_pvc_name - ) - if init_container: - if "initContainers" not in pod_spec["spec"]: - pod_spec["spec"]["initContainers"] = [] - pod_spec["spec"]["initContainers"].append(init_container) - - blocking_type = task.get("blocking", "completion") - if blocking_type == "endpoint": - service = create_service(task, task_id, job_id, jobset_name, namespace) - if service: - services.append(service) - - replicated_job = { - "name": f"task-{task_id}", - "replicas": 1, - "template": { - "spec": { - "parallelism": 1, - "completions": 1, - "backoffLimit": 0, - "template": pod_spec, - } - }, - } - - if previous_job_name: - dependency_status = ( - "Ready" if previous_blocking_type == "endpoint" else "Complete" - ) - replicated_job["dependsOn"] = [ - {"name": previous_job_name, "status": dependency_status} - ] - - replicated_jobs.append(replicated_job) - previous_job_name = replicated_job["name"] - previous_blocking_type = blocking_type - - return replicated_jobs, services - - -def generate_jobset(job_spec, node_assignments): - """ - Generate JobSet YAML from job specification. - - Args: - job_spec: Job specification dict with name, namespace, tasks, storage - node_assignments: Dict mapping task index to node ID - - Returns: - tuple: (JobSet manifest, services list, base_dir) - """ - job_id = job_spec["name"] - jobset_name = f"coldpress-{job_id}" - namespace = job_spec["namespace"] - tasks = job_spec.get("tasks", []) - storage = job_spec.get("storage", {}) - - # Extract discovery configuration - discovery_config = job_spec.get("discovery", {}) - discovery_template = discovery_config.get("template") - discovery_tasks = discovery_config.get("tasks", "all") - - # Determine which task indices should run discovery - if discovery_template: - if discovery_tasks == "all": - discovery_task_indices = set(range(len(tasks))) - else: - discovery_task_indices = set(discovery_tasks) - else: - discovery_task_indices = set() - - data_pvc_name = storage.get("results", f"coldpress-{namespace}-storage") - model_pvc_name = storage.get("models", "coldpress-model-storage") - base_dir = generate_base_dir(namespace, job_id) - - # Preprocess tasks - for task_id, task in enumerate(tasks): - _infer_blocking_type_and_health_check(task, task_id, job_id, namespace) - _substitute_dns_in_args(tasks, job_id, namespace) - - # Build init jobs (just mkdir, discovery now runs as init containers per task) - init_jobs, previous_job_name, previous_blocking_type = _build_init_jobs( - base_dir, data_pvc_name, len(tasks) - ) - - # Build task jobs - task_jobs, services = _build_task_jobs( - tasks, - job_id, - jobset_name, - namespace, - job_spec, - node_assignments, - data_pvc_name, - model_pvc_name, - base_dir, - previous_job_name, - previous_blocking_type, - discovery_template, - discovery_task_indices, - ) - - replicated_jobs = init_jobs + task_jobs - - # Determine driver jobs - driver_jobs = [ - f"task-{i}" - for i, task in enumerate(tasks) - if task.get("blocking", "completion") == "completion" - ] - - # Build JobSet spec - jobset_spec = {"suspend": True, "replicatedJobs": replicated_jobs} - if driver_jobs: - jobset_spec["successPolicy"] = { - "operator": "All", - "targetReplicatedJobs": driver_jobs, - } - - # Build JobSet manifest with standard labels - jobset_labels = COLDPRESS_LABELS.copy() - jobset_labels.update( - { - "kueue.x-k8s.io/queue-name": f"coldpress-local-queue-{namespace}", - "coldpress.io/job-id": jobset_name, - } - ) - - jobset = { - "apiVersion": "jobset.x-k8s.io/v1alpha2", - "kind": "JobSet", - "metadata": { - "name": jobset_name, - "namespace": namespace, - "labels": jobset_labels, - "annotations": { - "coldpress.io/base-dir": base_dir, - "coldpress.io/storage-pvc": data_pvc_name, - }, - }, - "spec": jobset_spec, - } - - return jobset, services, base_dir - - -def _extract_container_config(task): - """Extract container configuration from task.""" - containers = task.get("containers", []) - if containers: - container_def = containers[0] - return { - "image": container_def.get("image", "alpine:latest"), - "command": container_def.get("command"), - "args": container_def.get("args", []), - "env": container_def.get("env", []), - "working_dir": container_def.get("workingDir"), - "resources": container_def.get("resources", {}), - "gpus": int( - container_def.get("resources", {}) - .get("requests", {}) - .get("nvidia.com/gpu", "0") - ), - } - - # Fallback to old flat format - return { - "image": task.get("image", "alpine:latest"), - "command": task.get("command"), - "args": task.get("args", []), - "env": task.get("env", []), - "working_dir": task.get("workingDir"), - "resources": task.get("resources", {}), - "gpus": task.get("gpus", 0), - } - - -def _build_container_resources(resources, gpus): - """Build resource requests and limits.""" - if not resources.get("limits"): - resources["limits"] = resources.get("requests", {}).copy() - if not resources.get("requests"): - resources["requests"] = {} - - if gpus > 0: - resources["limits"]["nvidia.com/gpu"] = str(gpus) - resources["requests"]["nvidia.com/gpu"] = str(gpus) - - return { - "requests": resources.get("requests", {}), - "limits": resources.get("limits", {}), - } - - -def _build_volume_mounts(task, task_id, base_dir): - """Build volume mounts for container.""" - volume_mounts = [{"name": "coldpress-data", "mountPath": "/mnt/coldpress-data"}] - - result_path = task.get("result_path", f"{base_dir}/{task_id}") - ephemeral_mounts = task.get("ephemeral_mounts", []) - - if isinstance(ephemeral_mounts, str): - ephemeral_mounts = [ephemeral_mounts] - - for mount in ephemeral_mounts: - target = ( - mount.get("target", "/tmp/result") if isinstance(mount, dict) else mount - ) - volume_mounts.append( - { - "name": "coldpress-data", - "mountPath": target, - "subPath": result_path, - } - ) - - return volume_mounts - - -def _build_readiness_probe(task): - """Build readiness probe for endpoint blocking.""" - if task.get("blocking") != "endpoint" or not task.get("health_check"): - return None - - health_url = task["health_check"] - parsed = urlparse(health_url) - return { - "httpGet": { - "path": parsed.path or "/", - "port": parsed.port or 8000, - "scheme": (parsed.scheme or "http").upper(), - }, - "initialDelaySeconds": 30, - "periodSeconds": 10, - "failureThreshold": 10, - } - - -def build_container_spec( - task, task_id, job_id, namespace, data_pvc_name, model_pvc_name, base_dir -): - """Build container specification from task.""" - config = _extract_container_config(task) - - container = { - "name": "main", - "image": config["image"], - "volumeMounts": _build_volume_mounts(task, task_id, base_dir), - "resources": _build_container_resources(config["resources"], config["gpus"]), - } - - if config["working_dir"]: - container["workingDir"] = config["working_dir"] - if config["command"]: - container["command"] = ( - config["command"] - if isinstance(config["command"], list) - else shlex.split(config["command"]) - ) - if config["args"]: - container["args"] = ( - config["args"] - if isinstance(config["args"], list) - else shlex.split(config["args"]) - ) - if config["env"]: - container["env"] = ( - [{"name": k, "value": str(v)} for k, v in config["env"].items()] - if isinstance(config["env"], dict) - else config["env"] - ) - - readiness_probe = _build_readiness_probe(task) - if readiness_probe: - container["readinessProbe"] = readiness_probe - - return container - - -def _add_configmap_volume(volumes, container_spec, configmap_info): - """Add ConfigMap volume and mounts if specified.""" - if not configmap_info: - return - - volumes.append( - {"name": "configmap-files", "configMap": {"name": configmap_info["name"]}} - ) - for file_name in configmap_info.get("files", []): - container_spec["volumeMounts"].append( - { - "name": "configmap-files", - "mountPath": f"/workspace/{file_name}", - "subPath": file_name, - } - ) - - -def _add_task_volumes(volumes, container_spec, task, task_id, base_dir): - """Add volumes from task.volumes array.""" - for vol in task.get("volumes", []): - vol_name = vol.get("name") - vol_type = vol.get("type", "pvc") - mount_path = vol.get("mount") - - if vol_type == "emptyDir": - volume_def = {"name": vol_name, "emptyDir": {}} - if vol.get("medium"): - volume_def["emptyDir"]["medium"] = vol["medium"] - if vol.get("sizeLimit"): - volume_def["emptyDir"]["sizeLimit"] = vol["sizeLimit"] - volumes.append(volume_def) - - if mount_path: - if vol_name == "results": - # Mount results to task-specific subdirectory - task_subpath = ( - f"{base_dir}/task-{task_id}" if base_dir else vol.get("subPath", "") - ) - container_spec["volumeMounts"].append( - { - "name": "coldpress-data", - "mountPath": mount_path, - "subPath": task_subpath, - } - ) - else: - container_spec["volumeMounts"].append( - {"name": vol_name, "mountPath": mount_path} - ) - - -def _add_sys_mounts(volumes, container_spec, task): - """Add host path mounts if specified.""" - for i, mount in enumerate(task.get("sys_mounts", [])): - volumes.append( - { - "name": f"sys-{i}", - "hostPath": {"path": mount["source"], "type": "Directory"}, - } - ) - container_spec["volumeMounts"].append( - { - "name": f"sys-{i}", - "mountPath": mount["target"], - "readOnly": mount.get("read_only", False), - } - ) - - -def _apply_pod_options(pod_spec, task, container_spec): - """Apply optional pod settings like tolerations, network mode, and security context.""" - if task.get("tolerate_all"): - pod_spec["spec"]["tolerations"] = [{"operator": "Exists"}] - - if task.get("network_mode") == "host": - pod_spec["spec"]["hostNetwork"] = True - - if task.get("privileged"): - container_spec["securityContext"] = {"privileged": True} - - -def build_pod_spec( - task, - task_id, - job_id, - jobset_name, - node_id, - container_spec, - data_pvc_name, - configmap_info=None, - base_dir=None, -): - """Build pod specification. - - Args: - job_id: Base job name for DNS - jobset_name: Full JobSet name with prefix for labels - """ - volumes = [ - { - "name": "coldpress-data", - "persistentVolumeClaim": {"claimName": data_pvc_name}, - } - ] - - _add_configmap_volume(volumes, container_spec, configmap_info) - _add_task_volumes(volumes, container_spec, task, task_id, base_dir) - _add_sys_mounts(volumes, container_spec, task) - - pod_spec = { - "metadata": { - "labels": { - "app": f"task-{task_id}", - "coldpress/gid": jobset_name, - }, - "annotations": task.get("annotations", {}), - }, - "spec": { - "restartPolicy": "Never", - "volumes": volumes, - "containers": [container_spec], - }, - } - - # Add node selection - either specific node or any coldpress node - if node_id is not None and node_id != "any": - # Pin to specific node - pod_spec["spec"]["nodeSelector"] = {"coldpress.node": node_id} - else: - # Let Kubernetes scheduler pick any node with coldpress.node label - pod_spec["spec"]["affinity"] = { - "nodeAffinity": { - "requiredDuringSchedulingIgnoredDuringExecution": { - "nodeSelectorTerms": [ - { - "matchExpressions": [ - { - "key": "coldpress.node", - "operator": "Exists", - } - ] - } - ] - } - } - } - - _apply_pod_options(pod_spec, task, container_spec) - - return pod_spec - - -def create_service(task, task_id, job_id, jobset_name, namespace): - """Create Kubernetes Service for endpoint blocking. - - Args: - job_id: Base job name for service naming (e.g., "ddp-training") - jobset_name: Full JobSet name with prefix for labels (e.g., "coldpress-ddp-training") - """ - health_check = task.get("health_check") - if not health_check: - return None - - try: - parsed = urlparse(health_check) - port = parsed.port or 8000 - - # Merge standard labels with job-specific labels - service_labels = COLDPRESS_LABELS.copy() - service_labels.update( - { - "coldpress/gid": jobset_name, - "coldpress.io/job-id": jobset_name, - } - ) - - service = { - "apiVersion": "v1", - "kind": "Service", - "metadata": { - "name": f"coldpress-s-{job_id}-{task_id}", - "namespace": namespace, - "labels": service_labels, - }, - "spec": { - "selector": { - "app": f"task-{task_id}", - "coldpress/gid": jobset_name, - }, - "ports": [{"port": port, "targetPort": port}], - "type": "ClusterIP", - }, - } - return service - except (ValueError, AttributeError, KeyError) as e: - # URL parsing failed or service configuration is invalid - sys.stderr.write(f"Warning: Could not create service for task {task_id}: {e}\n") - return None - - -def build_discovery_init_container(template_path, task_id, base_dir, data_pvc_name): - """ - Build discovery init container from template. - - Args: - template_path: Path to discovery template YAML - task_id: Task ID for result path - base_dir: Base directory path - data_pvc_name: PVC name - - Returns: - dict: Init container spec, or None if template not found - """ - try: - # Read discovery template - with open(template_path, "r") as f: - template = yaml.safe_load(f) - - # Extract Pod spec - pod_spec = template.get("spec", {}) - containers = pod_spec.get("containers", []) - - if not containers: - return None - - # Get template name from filename - template_name = ( - os.path.basename(template_path).replace(".yaml", "").replace(".yml", "") - ) - - # Copy container and modify for init container use - container = containers[0].copy() - container["name"] = "discovery" - - # Task-specific output path - task_result_path = f"{base_dir}/task-{task_id}" - - # Update volume mounts to use PVC with task-specific path - container["volumeMounts"] = [ - { - "name": "coldpress-data", - "mountPath": "/tmp/result", - "subPath": task_result_path, - } - ] - - # Add rename command to output discovery_{template_name}.json - original_command = ( - container.get("args", [""])[0] if container.get("args") else "" - ) - rename_cmd = f"\nif [ -f /tmp/result/discovery.json ]; then mv /tmp/result/discovery.json /tmp/result/discovery_{template_name}.json; fi" - - if container.get("args"): - container["args"] = [original_command + rename_cmd] - - return container - except (FileNotFoundError, yaml.YAMLError, KeyError, IndexError) as e: - sys.stderr.write( - f"Warning: Could not load discovery template {template_path}: {e}\n" - ) - return None - - -def build_mkdir_job(base_dir, data_pvc_name, namespace, num_tasks): - """ - Build mkdir init job to create base directory and task subdirectories. - - Args: - base_dir: Base directory path - data_pvc_name: PVC name - namespace: Kubernetes namespace - num_tasks: Number of tasks (to create task-N directories) - - Returns: - dict: ReplicatedJob for mkdir - """ - # Build command to create base dir and all task subdirectories - task_dirs = " ".join([f"/data/{base_dir}/task-{i}" for i in range(num_tasks)]) - mkdir_cmd = f"mkdir -p /data/{base_dir} {task_dirs} && echo Created directory /data/{base_dir} with {num_tasks} task subdirectories" - - return { - "name": "mkdir", - "replicas": 1, - "template": { - "spec": { - "parallelism": 1, - "completions": 1, - "backoffLimit": 0, - "template": { - "metadata": {"labels": {"app": "mkdir"}}, - "spec": { - "restartPolicy": "Never", - "containers": [ - { - "name": "mkdir", - "image": "registry.access.redhat.com/ubi9/ubi-minimal:latest", - "command": ["sh", "-c", mkdir_cmd], - "volumeMounts": [ - {"name": "storage", "mountPath": "/data"} - ], - } - ], - "volumes": [ - { - "name": "storage", - "persistentVolumeClaim": {"claimName": data_pvc_name}, - } - ], - }, - }, - } - }, - } - - -def build_discovery_job(template_path, base_dir, data_pvc_name, node_id): - """ - Build discovery job from template. - - Args: - template_path: Path to discovery template YAML - base_dir: Base directory path - data_pvc_name: PVC name - node_id: Node ID to run on - - Returns: - dict: ReplicatedJob for discovery, or None if template not found - """ - try: - # Read discovery template - with open(template_path, "r") as f: - template = yaml.safe_load(f) - - # Extract Pod spec - pod_spec = template.get("spec", {}) - containers = pod_spec.get("containers", []) - - if not containers: - return None - - # Get template name from filename - template_name = ( - os.path.basename(template_path).replace(".yaml", "").replace(".yml", "") - ) - - # Modify container to write to base_dir - container = containers[0].copy() - - # Update volume mounts to use PVC - container["volumeMounts"] = [ - {"name": "storage", "mountPath": "/tmp/result", "subPath": base_dir} - ] - - # Add rename command to output discovery_{template_name}.json - original_command = ( - container.get("args", [""])[0] if container.get("args") else "" - ) - rename_cmd = f"\nif [ -f /tmp/result/discovery.json ]; then mv /tmp/result/discovery.json /tmp/result/discovery_{template_name}.json; fi" - - if container.get("args"): - container["args"] = [original_command + rename_cmd] - - # Build pod spec with node selection - pod_spec = { - "restartPolicy": "Never", - "containers": [container], - "volumes": [ - { - "name": "storage", - "persistentVolumeClaim": {"claimName": data_pvc_name}, - } - ], - "tolerations": [{"operator": "Exists"}], - } - - # Add node selection - either specific node or any coldpress node - if node_id is not None and node_id != "any": - # Pin to specific node - pod_spec["nodeSelector"] = {"coldpress.node": str(node_id)} - else: - # Let Kubernetes scheduler pick any node with coldpress.node label - pod_spec["affinity"] = { - "nodeAffinity": { - "requiredDuringSchedulingIgnoredDuringExecution": { - "nodeSelectorTerms": [ - { - "matchExpressions": [ - { - "key": "coldpress.node", - "operator": "Exists", - } - ] - } - ] - } - } - } - - return { - "name": "discovery", - "replicas": 1, - "template": { - "spec": { - "parallelism": 1, - "completions": 1, - "backoffLimit": 0, - "template": { - "metadata": {"labels": {"app": "discovery"}}, - "spec": pod_spec, - }, - } - }, - } - except (FileNotFoundError, yaml.YAMLError, KeyError, IndexError) as e: - sys.stderr.write( - f"Warning: Could not load discovery template {template_path}: {e}\n" - ) - return None - - -def jobset_to_yaml(jobset): - """Convert JobSet dict to YAML string.""" - return yaml.dump(jobset, default_flow_style=False, sort_keys=False) - - -def services_to_yaml(services): - """Convert list of services to YAML string.""" - if not services: - return "" - return yaml.dump_all(services, default_flow_style=False, sort_keys=False) diff --git a/coldpress/jobset_generator.py b/coldpress/jobset_generator.py new file mode 100644 index 0000000..0018541 --- /dev/null +++ b/coldpress/jobset_generator.py @@ -0,0 +1,456 @@ +# Assisted by: Claude Sonnet 4.5 +"""JobSet generation from vanilla k8s Jobs + intent.yaml. + +Generates JobSet manifests directly from: +- Vanilla Kubernetes Job manifests (job-spec.yaml) +- Intent specifications (intent.yaml) + +Direct transformation from vk8s to JobSet with no intermediate formats. +""" + +import copy +import yaml +from datetime import datetime, timezone +from .constants import ( + COLDPRESS_LABELS, + MKDIR_IMAGE, + get_kueue_queue_label, + get_jobset_name, +) +from .utils import ( + substitute_macros, + apply_arg_overrides, + apply_env_overrides, + build_discovery_init_container, + parse_discovery_config, + get_storage_pvc_name, + DANGEROUS_SHELL_CHARS, +) + + +def build_replica_macros( + task_name: str, + replica_index: int, + total_replicas: int, + node_id: int | None, + jobset_name: str, + namespace: str, +) -> dict[str, str]: + """Build macro substitution dict for a replica.""" + macros = { + "INDEX": str(replica_index), + "REPLICAS": str(total_replicas), + "TASK_NAME": task_name, + } + + if node_id is not None: + macros["NODE_ID"] = str(node_id) + + # REPLICA_N macros - full JobSet DNS + for i in range(total_replicas): + macros[f"REPLICA_{i}"] = ( + f"{jobset_name}-{task_name}-{i}-0.{jobset_name}-{task_name}.{namespace}.svc.cluster.local" + ) + + # REPLICAS_ALL + all_replicas = ",".join( + f"{jobset_name}-{task_name}-{i}-0" for i in range(total_replicas) + ) + macros["REPLICAS_ALL"] = all_replicas + + return macros + + +def generate_jobset_from_intent( + vk8s_jobs: dict, intent_config, project_config, namespace: str +): + """ + Generate JobSet manifest from vanilla k8s Jobs + intent. + + Args: + vk8s_jobs: Dict mapping job name to vanilla k8s Job manifest + intent_config: Validated IntentConfig object + project_config: Project configuration dict + namespace: Target namespace + + Returns: + tuple: (JobSet manifest, services list, base_dir) + """ + # Determine job/jobset name + if len(intent_config.tasks) == 1: + job_id = intent_config.tasks[0].name + else: + job_id = f"{intent_config.tasks[0].name}-workflow" + + jobset_name = get_jobset_name(job_id) + data_pvc_name = get_storage_pvc_name(project_config, namespace) + + # Generate base directory for results + timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") + short_hash = abs(hash(jobset_name + timestamp)) % 100000000 + base_dir = f"{namespace}/coldpress_results/{job_id}-{short_hash:08x}-{timestamp}" + + # Extract discovery configuration + discovery_template_path, discovery_task_names = parse_discovery_config( + intent_config + ) + + # Build replicated jobs + replicated_jobs = [] + task_counter = 0 + + # Pre-scan tasks to build task-qualified REPLICA macros + # Maps REPLICA_taskname_N -> pod DNS for all tasks + task_replica_macros = {} + temp_task_counter = 0 + + for task_intent in intent_config.tasks: + task_name = task_intent.name + replicas = task_intent.replicas or 1 + + # Build task-qualified REPLICA macros for this task + for i in range(replicas): + replica_job_name = f"task-{temp_task_counter + i}" + pod_dns = f"{jobset_name}-{replica_job_name}-0-0.{jobset_name}.{namespace}.svc.cluster.local" + task_replica_macros[f"REPLICA_{task_name}_{i}"] = pod_dns + + temp_task_counter += replicas + + # Create mkdir init job + task_count = sum(task.replicas or 1 for task in intent_config.tasks) + mkdir_job = build_mkdir_job(base_dir, data_pvc_name, task_count, namespace) + replicated_jobs.append(mkdir_job) + + previous_job_name = "mkdir" + previous_blocking_type = "completion" + + for task_intent in intent_config.tasks: + task_name = task_intent.name + + if task_name not in vk8s_jobs: + raise ValueError(f"Task '{task_name}' not found in job-spec.yaml") + + vk8s_job = vk8s_jobs[task_name] + replicas = task_intent.replicas or 1 + + # For multi-replica tasks (like DDP), all replicas should depend on mkdir + # and run in parallel, not sequentially + task_dependency_job = previous_job_name + task_dependency_type = previous_blocking_type + + # Build list of replicated job names for this task + replica_job_names = [f"task-{task_counter + i}" for i in range(replicas)] + + for replica_idx in range(replicas): + # Build macros for this replica using actual replicated job names + nodes = task_intent.nodes + node_id = nodes[replica_idx] if nodes else None + + # Build macros with actual JobSet DNS names + macros = { + "INDEX": str(replica_idx), + "REPLICAS": str(replicas), + "TASK_NAME": task_name, + } + + if node_id is not None: + macros["NODE_ID"] = str(node_id) + + # REPLICA_N macros - use actual replicated job names + for i in range(replicas): + replica_job_name = replica_job_names[i] + # JobSet DNS format: {jobset-name}-{replicated-job-name}-{job-index}-{completion-index}.{subdomain}.{namespace}.svc.cluster.local + macros[f"REPLICA_{i}"] = ( + f"{jobset_name}-{replica_job_name}-0-0.{jobset_name}.{namespace}.svc.cluster.local" + ) + + # REPLICAS_ALL + all_replicas = ",".join( + f"{jobset_name}-{rjn}-0-0" for rjn in replica_job_names + ) + macros["REPLICAS_ALL"] = all_replicas + + # Add task-qualified REPLICA macros for all tasks + macros.update(task_replica_macros) + + # Determine if this task should run discovery + should_run_discovery = ( + discovery_template_path and task_name in discovery_task_names + ) + + # Transform vk8s Job to JobSet replicated job + # All replicas of the same task share the same dependency + replicated_job = build_replicated_job_from_vk8s( + vk8s_job, + task_intent, + replica_idx, + macros, + task_counter, + jobset_name, + namespace, + data_pvc_name, + base_dir, + task_dependency_job, # All replicas depend on the same previous job + task_dependency_type, + discovery_template_path if should_run_discovery else None, + ) + + replicated_jobs.append(replicated_job) + + task_counter += 1 + + # After all replicas of this task, update previous for next task + # Next task depends on the last replica of current task + previous_job_name = replicated_jobs[-1]["name"] + # Determine blocking type from readiness probe + vk8s_spec = vk8s_job.get("spec", {}).get("template", {}).get("spec", {}) + containers = vk8s_spec.get("containers", []) + has_readiness = containers and containers[0].get("readinessProbe") + previous_blocking_type = "endpoint" if has_readiness else "completion" + + # Determine driver jobs (non-endpoint tasks) + driver_jobs = [ + job["name"] + for job in replicated_jobs + if job["name"] != "mkdir" # Skip mkdir + ] + + # Build JobSet spec + jobset_spec = { + "suspend": True, + "replicatedJobs": replicated_jobs, + } + + if driver_jobs: + jobset_spec["successPolicy"] = { + "operator": "All", + "targetReplicatedJobs": driver_jobs, + } + + # Build JobSet manifest + jobset_labels = COLDPRESS_LABELS.copy() + jobset_labels.update( + { + "kueue.x-k8s.io/queue-name": get_kueue_queue_label(namespace), + "coldpress.io/job-id": jobset_name, + } + ) + + jobset = { + "apiVersion": "jobset.x-k8s.io/v1alpha2", + "kind": "JobSet", + "metadata": { + "name": jobset_name, + "namespace": namespace, + "labels": jobset_labels, + "annotations": { + "coldpress.io/base-dir": base_dir, + "coldpress.io/storage-pvc": data_pvc_name, + }, + }, + "spec": jobset_spec, + } + + # Services are no longer generated - JobSet provides automatic DNS + return jobset, [], base_dir + + +def build_mkdir_job( + base_dir: str, pvc_name: str, task_count: int, namespace: str +) -> dict: + """Build mkdir initialization job.""" + # Validate base_dir contains no shell metacharacters (defense in depth) + for char in DANGEROUS_SHELL_CHARS: + if char in base_dir: + raise ValueError( + f"Invalid character '{char}' in base_dir '{base_dir}'. " + f"This could be a security risk." + ) + + # SAFETY: base_dir validated above, task_count is int - safe for shell command + # Create task subdirectories + task_dirs = " ".join(f"/data/{base_dir}/task-{i}" for i in range(task_count)) + + return { + "name": "mkdir", + "replicas": 1, + "template": { + "spec": { + "parallelism": 1, + "completions": 1, + "backoffLimit": 0, + "template": { + "metadata": {"labels": {"app": "mkdir"}}, + "spec": { + "restartPolicy": "Never", + "containers": [ + { + "name": "mkdir", + "image": MKDIR_IMAGE, + "command": ["sh", "-c"], + "args": [ + f"mkdir -p /data/{base_dir} {task_dirs} && " + f"echo Created directory /data/{base_dir} with {task_count} task subdirectories" + ], + "volumeMounts": [ + {"name": "storage", "mountPath": "/data"} + ], + } + ], + "volumes": [ + { + "name": "storage", + "persistentVolumeClaim": {"claimName": pvc_name}, + } + ], + }, + }, + }, + }, + } + + +def build_replicated_job_from_vk8s( + vk8s_job: dict, + task_intent, + replica_idx: int, + macros: dict, + task_counter: int, + jobset_name: str, + namespace: str, + pvc_name: str, + base_dir: str, + previous_job_name: str, + previous_blocking_type: str, + discovery_template_path: str = None, +): + """Build a JobSet replicated job from a vanilla k8s Job.""" + # Extract vk8s Job spec + vk8s_spec = vk8s_job.get("spec", {}).get("template", {}).get("spec", {}) + containers = vk8s_spec.get("containers", []) + + if not containers: + raise ValueError( + f"Job {vk8s_job.get('metadata', {}).get('name')} has no containers" + ) + + container = copy.deepcopy(containers[0]) + + # Substitute macros in existing env vars from job-spec + if container.get("env"): + for env_item in container["env"]: + if "value" in env_item and isinstance(env_item["value"], str): + context = f" for env var '{env_item['name']}'" + env_item["value"] = substitute_macros( + env_item["value"], macros, context + ) + + # Apply arg and env var replacements using shared utilities + apply_arg_overrides(container, task_intent, macros) + apply_env_overrides(container, task_intent, macros) + + # Build pod spec + pod_spec = { + "restartPolicy": "Never", + "containers": [container], + } + + # Copy tolerations + if vk8s_spec.get("tolerations"): + pod_spec["tolerations"] = vk8s_spec["tolerations"] + + # Copy volumes and add coldpress- prefix to configMap names + if vk8s_spec.get("volumes"): + pod_spec["volumes"] = copy.deepcopy(vk8s_spec["volumes"]) + # Add coldpress- prefix to configMap volume names + for volume in pod_spec["volumes"]: + if "configMap" in volume: + cm_name = volume["configMap"]["name"] + if not cm_name.startswith("coldpress-"): + volume["configMap"]["name"] = f"coldpress-{cm_name}" + + # Find storage volume name for discovery init container + storage_volume_name = None + if vk8s_spec.get("volumes"): + for volume in vk8s_spec["volumes"]: + if "persistentVolumeClaim" in volume: + pvc_claim = volume["persistentVolumeClaim"].get("claimName") + if pvc_claim == pvc_name: + storage_volume_name = volume.get("name") + break + + # Add discovery init container if specified + if discovery_template_path and storage_volume_name: + discovery_init = build_discovery_init_container( + discovery_template_path, base_dir, storage_volume_name, task_counter + ) + if discovery_init: + pod_spec["initContainers"] = [discovery_init] + + # Transform volumeMounts to add subPath for PVC mounts pointing to task directory + if container.get("volumeMounts"): + volume_mounts = copy.deepcopy(container["volumeMounts"]) + task_subdir = f"{base_dir}/task-{task_counter}" + + # Find which volumes are PVCs using the storage PVC + pvc_volume_names = [] + if vk8s_spec.get("volumes"): + for volume in vk8s_spec["volumes"]: + if "persistentVolumeClaim" in volume: + pvc_claim = volume["persistentVolumeClaim"].get("claimName") + if pvc_claim == pvc_name: + pvc_volume_names.append(volume.get("name")) + + # Add subPath to volumeMounts that reference the storage PVC + for mount in volume_mounts: + if mount.get("name") in pvc_volume_names and not mount.get("subPath"): + # Add subPath to mount at task-specific directory created by mkdir + mount["subPath"] = task_subdir + + container["volumeMounts"] = volume_mounts + + # Build replicated job + replicated_job = { + "name": f"task-{task_counter}", + "replicas": 1, + "template": { + "spec": { + "parallelism": 1, + "completions": 1, + "backoffLimit": 0, + "template": { + "metadata": { + "labels": { + "app": f"task-{task_counter}", + "coldpress/gid": jobset_name, + }, + }, + "spec": pod_spec, + }, + }, + }, + } + + # Add dependency + if previous_job_name: + dependency_status = ( + "Ready" if previous_blocking_type == "endpoint" else "Complete" + ) + replicated_job["dependsOn"] = [ + {"name": previous_job_name, "status": dependency_status} + ] + + # JobSet provides automatic DNS service - no need to create additional services + return replicated_job + + +def jobset_to_yaml(jobset: dict) -> str: + """Convert JobSet dict to YAML string.""" + + return yaml.safe_dump(jobset, default_flow_style=False, sort_keys=False) + + +def services_to_yaml(services: list) -> str: + """Convert Services list to YAML string.""" + + return yaml.safe_dump_all(services, default_flow_style=False, sort_keys=False) diff --git a/coldpress/kubeflow_generator.py b/coldpress/kubeflow_generator.py new file mode 100644 index 0000000..6fe2f0a --- /dev/null +++ b/coldpress/kubeflow_generator.py @@ -0,0 +1,370 @@ +# Assisted by: Claude Sonnet 4.5 +"""Kubeflow manifest generation for Coldpress jobs.""" + +import copy +import yaml +from datetime import datetime, timezone + +from .constants import ( + COLDPRESS_LABELS, + MKDIR_IMAGE, + get_kueue_queue_label, +) +from .utils import ( + apply_arg_overrides, + apply_env_overrides, + build_discovery_init_container, + parse_discovery_config, + get_storage_pvc_name, + DANGEROUS_SHELL_CHARS, +) + + +def generate_inferenceservice_from_intent( + vk8s_jobs: dict, intent_config, project_config, namespace: str +): + """ + Generate KServe InferenceService manifest from vanilla k8s Jobs + intent. + + Args: + vk8s_jobs: Dict mapping job name to vanilla k8s Job manifest + intent_config: Validated IntentConfig object + project_config: Project configuration dict + namespace: Target namespace + + Returns: + tuple: (InferenceService manifest, base_dir) + """ + # Only support single task for KServe + if len(intent_config.tasks) != 1: + raise ValueError("KServe generation currently only supports single task") + + task_intent = intent_config.tasks[0] + task_name = task_intent.name + + if task_name not in vk8s_jobs: + raise ValueError(f"Task '{task_name}' not found in job-spec.yaml") + + vk8s_job = vk8s_jobs[task_name] + + # Determine job name + job_id = task_name + inferenceservice_name = f"coldpress-{job_id}" + + data_pvc_name = get_storage_pvc_name(project_config, namespace) + + # Generate base directory for results/logs + timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") + short_hash = abs(hash(inferenceservice_name + timestamp)) % 100000000 + base_dir = f"{namespace}/coldpress_results/{job_id}-{short_hash:08x}-{timestamp}" + + # Extract vk8s Job spec + vk8s_spec = vk8s_job.get("spec", {}).get("template", {}).get("spec", {}) + containers = vk8s_spec.get("containers", []) + + if not containers: + raise ValueError( + f"Job {vk8s_job.get('metadata', {}).get('name')} has no containers" + ) + + container = copy.deepcopy(containers[0]) + + # KServe requires specific container name + container["name"] = "kserve-container" + + # Get replicas (for autoscaling configuration) + replicas = task_intent.replicas or 1 + + # KServe macros + macros = { + "REPLICAS": str(replicas), + "TASK_NAME": task_name, + } + + # Apply arg and env var replacements using shared utilities + apply_arg_overrides(container, task_intent, macros) + apply_env_overrides(container, task_intent, macros) + + # Extract ports from container spec + container.get("ports", []) + + # Build predictor spec with custom container + predictor = {"containers": [container]} + + # Copy volumes if present + if vk8s_spec.get("volumes"): + predictor["volumes"] = copy.deepcopy(vk8s_spec["volumes"]) + # Add coldpress- prefix to configMap volume names + for volume in predictor["volumes"]: + if "configMap" in volume: + cm_name = volume["configMap"]["name"] + if not cm_name.startswith("coldpress-"): + volume["configMap"]["name"] = f"coldpress-{cm_name}" + + # Copy tolerations if present + if vk8s_spec.get("tolerations"): + predictor["tolerations"] = vk8s_spec["tolerations"] + + # Add node selector if specified + node_id = task_intent.nodes[0] if task_intent.nodes else None + if node_id and node_id != "any": + predictor["nodeSelector"] = {"coldpress.node": str(node_id)} + + # Build InferenceService labels + inferenceservice_labels = COLDPRESS_LABELS.copy() + inferenceservice_labels.update( + { + "coldpress.io/job-id": inferenceservice_name, + "kueue.x-k8s.io/queue-name": get_kueue_queue_label(namespace), + } + ) + + # Build InferenceService annotations + annotations = { + "coldpress.io/base-dir": base_dir, + "coldpress.io/storage-pvc": data_pvc_name, + "serving.kserve.io/deploymentMode": "RawDeployment", # Use raw K8s deployment for more control + } + + # Add autoscaling annotations if replicas specified + if replicas > 1: + annotations["serving.kserve.io/autoscalerClass"] = "hpa" + annotations["serving.kserve.io/minReplicas"] = "1" + annotations["serving.kserve.io/maxReplicas"] = str(replicas) + + # Build InferenceService manifest + inferenceservice = { + "apiVersion": "serving.kserve.io/v1beta1", + "kind": "InferenceService", + "metadata": { + "name": inferenceservice_name, + "namespace": namespace, + "labels": inferenceservice_labels, + "annotations": annotations, + }, + "spec": {"predictor": predictor}, + } + + return inferenceservice, base_dir + + +def kubeflow_to_yaml(manifest): + """Convert Kubeflow manifest dict to YAML string. + + Args: + manifest: Kubeflow manifest dict + + Returns: + str: YAML string + """ + return yaml.safe_dump(manifest, default_flow_style=False, sort_keys=False) + + +def generate_pytorchjob_from_intent( + vk8s_jobs: dict, intent_config, project_config, namespace: str +): + """ + Generate PyTorchJob manifest from vanilla k8s Jobs + intent. + + Args: + vk8s_jobs: Dict mapping job name to vanilla k8s Job manifest + intent_config: Validated IntentConfig object + project_config: Project configuration dict + namespace: Target namespace + + Returns: + tuple: (PyTorchJob manifest, base_dir) + """ + # Only support single task for Kubeflow + if len(intent_config.tasks) != 1: + raise ValueError("Kubeflow generation currently only supports single task") + + task_intent = intent_config.tasks[0] + task_name = task_intent.name + + if task_name not in vk8s_jobs: + raise ValueError(f"Task '{task_name}' not found in job-spec.yaml") + + vk8s_job = vk8s_jobs[task_name] + + # Determine job name + job_id = task_name + pytorchjob_name = f"coldpress-{job_id}" + + data_pvc_name = get_storage_pvc_name(project_config, namespace) + + # Generate base directory for results + timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") + short_hash = abs(hash(pytorchjob_name + timestamp)) % 100000000 + base_dir = f"{namespace}/coldpress_results/{job_id}-{short_hash:08x}-{timestamp}" + + # Extract discovery configuration + discovery_template_path, discovery_task_names = parse_discovery_config( + intent_config + ) + + # Extract vk8s Job spec + vk8s_spec = vk8s_job.get("spec", {}).get("template", {}).get("spec", {}) + containers = vk8s_spec.get("containers", []) + + if not containers: + raise ValueError( + f"Job {vk8s_job.get('metadata', {}).get('name')} has no containers" + ) + + container = copy.deepcopy(containers[0]) + + # Kubeflow requires container to be named "pytorch" + container["name"] = "pytorch" + + # Get replicas (PyTorchJob workers) + replicas = task_intent.replicas or 1 + + # Kubeflow macros - simpler because PyTorchJob auto-injects env vars + macros = { + "REPLICAS": str(replicas), + "TASK_NAME": task_name, + } + + # Apply arg and env var replacements using shared utilities + apply_arg_overrides(container, task_intent, macros) + apply_env_overrides(container, task_intent, macros) + + # Add subPath to results volume mount in main container + if "volumeMounts" in container: + for volume_mount in container["volumeMounts"]: + if volume_mount.get("name") == "results": + volume_mount["subPath"] = base_dir + + # Build pod spec + pod_spec = { + "restartPolicy": "OnFailure", + "containers": [container], + } + + # Copy tolerations + if vk8s_spec.get("tolerations"): + pod_spec["tolerations"] = vk8s_spec["tolerations"] + + # Copy volumes and add coldpress- prefix to configMap names + if vk8s_spec.get("volumes"): + pod_spec["volumes"] = copy.deepcopy(vk8s_spec["volumes"]) + # Add coldpress- prefix to configMap volume names + for volume in pod_spec["volumes"]: + if "configMap" in volume: + cm_name = volume["configMap"]["name"] + if not cm_name.startswith("coldpress-"): + volume["configMap"]["name"] = f"coldpress-{cm_name}" + + # Add storage PVC if not already present + if not any(v.get("name") == "results" for v in pod_spec.get("volumes", [])): + if "volumes" not in pod_spec: + pod_spec["volumes"] = [] + pod_spec["volumes"].insert( + 0, + {"name": "results", "persistentVolumeClaim": {"claimName": data_pvc_name}}, + ) + + # Build init containers + init_containers = [] + + # Validate base_dir contains no shell metacharacters (defense in depth) + for char in DANGEROUS_SHELL_CHARS: + if char in base_dir: + raise ValueError( + f"Invalid character '{char}' in base_dir '{base_dir}'. " + f"This could be a security risk." + ) + + # Add mkdir init container + mkdir_init = { + "name": "mkdir", + "image": MKDIR_IMAGE, + "command": ["sh", "-c"], + "args": [ + f"mkdir -p /results/{base_dir} && echo Created directory /results/{base_dir}" + ], + "volumeMounts": [{"name": "results", "mountPath": "/results"}], + } + init_containers.append(mkdir_init) + + # Add discovery init container if specified + should_run_discovery = discovery_template_path and task_name in discovery_task_names + + if should_run_discovery: + # Find storage volume name + storage_volume_name = None + for volume in pod_spec.get("volumes", []): + if "persistentVolumeClaim" in volume: + pvc_claim = volume["persistentVolumeClaim"].get("claimName") + if pvc_claim == data_pvc_name: + storage_volume_name = volume.get("name") + break + + if storage_volume_name: + discovery_init = build_discovery_init_container( + discovery_template_path, + base_dir, + storage_volume_name, + task_id=None, + per_pod_directory=True, + ) + if discovery_init: + init_containers.append(discovery_init) + + pod_spec["initContainers"] = init_containers + + # Build pod template + pod_template = { + "metadata": { + "annotations": {"sidecar.istio.io/inject": "false"}, + "labels": COLDPRESS_LABELS.copy(), + }, + "spec": pod_spec, + } + pod_template["metadata"]["labels"]["coldpress.io/job-id"] = pytorchjob_name + + # Build PyTorchJob labels + pytorchjob_labels = COLDPRESS_LABELS.copy() + pytorchjob_labels.update( + { + "coldpress.io/job-id": pytorchjob_name, + "kueue.x-k8s.io/queue-name": get_kueue_queue_label(namespace), + } + ) + + # Build PyTorchJob manifest + pytorchjob = { + "apiVersion": "kubeflow.org/v1", + "kind": "PyTorchJob", + "metadata": { + "name": pytorchjob_name, + "namespace": namespace, + "labels": pytorchjob_labels, + "annotations": { + "coldpress.io/base-dir": base_dir, + "coldpress.io/storage-pvc": data_pvc_name, + "kueue.x-k8s.io/wait-for-pods-ready-timeout": "20m", + }, + }, + "spec": { + "suspend": True, # Start suspended for Kueue + "pytorchReplicaSpecs": { + "Master": { + "replicas": 1, + "restartPolicy": "OnFailure", + "template": pod_template, + } + }, + }, + } + + # Add Worker replicas if more than 1 replica + if replicas > 1: + pytorchjob["spec"]["pytorchReplicaSpecs"]["Worker"] = { + "replicas": replicas - 1, + "restartPolicy": "OnFailure", + "template": pod_template, + } + + return pytorchjob, base_dir diff --git a/coldpress/kuberay_generator.py b/coldpress/kuberay_generator.py new file mode 100644 index 0000000..6377189 --- /dev/null +++ b/coldpress/kuberay_generator.py @@ -0,0 +1,495 @@ +# Assisted by: Claude Sonnet 4.5 +"""KubeRay manifest generation for Coldpress jobs. + +Transforms single-node Ray Train jobs into distributed RayJob manifests. +""" + +import copy +import yaml +from datetime import datetime, timezone +from .constants import ( + COLDPRESS_LABELS, + MKDIR_IMAGE, + get_kueue_queue_label, +) +from .utils import ( + apply_arg_overrides, + apply_env_overrides, + build_discovery_init_container, + parse_discovery_config, + get_storage_pvc_name, + DANGEROUS_SHELL_CHARS, +) + + +def generate_rayjob_from_intent( + vk8s_jobs: dict, intent_config, project_config, namespace: str +): + """ + Generate RayJob manifest from vanilla k8s Job + intent. + + Takes a single-node Ray Train job and transforms it into a distributed RayJob. + + Note: For RayJob, the entrypoint runs via ray job submit which doesn't have + access to container mounts (ConfigMaps). Files must be in shared storage or + included in runtime environment. + + Args: + vk8s_jobs: Dict mapping job name to vanilla k8s Job manifest + intent_config: Validated IntentConfig object + project_config: Project configuration dict + namespace: Target namespace + + Returns: + tuple: (RayJob manifest, base_dir) + """ + # Only support single task for KubeRay + if len(intent_config.tasks) != 1: + raise ValueError("KubeRay generation currently only supports single task") + + task_intent = intent_config.tasks[0] + task_name = task_intent.name + + if task_name not in vk8s_jobs: + raise ValueError(f"Task '{task_name}' not found in job-spec.yaml") + + vk8s_job = vk8s_jobs[task_name] + + # Determine job name + job_id = task_name + rayjob_name = f"coldpress-{job_id}" + data_pvc_name = get_storage_pvc_name(project_config, namespace) + + # Generate base directory for results + timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") + short_hash = abs(hash(rayjob_name + timestamp)) % 100000000 + base_dir = f"{namespace}/coldpress_results/{job_id}-{short_hash:08x}-{timestamp}" + + # Extract discovery configuration + discovery_template_path, discovery_task_names = parse_discovery_config( + intent_config + ) + + # Extract vk8s Job spec (single-node pod template) + vk8s_spec = vk8s_job.get("spec", {}).get("template", {}).get("spec", {}) + containers = vk8s_spec.get("containers", []) + + if not containers: + raise ValueError( + f"Job {vk8s_job.get('metadata', {}).get('name')} has no containers" + ) + + base_container = copy.deepcopy(containers[0]) + + # Get replicas (number of Ray pods to create) + replicas = task_intent.replicas or 1 + if replicas < 1: + raise ValueError(f"replicas must be >= 1, got {replicas}") + + # Build macros for RayJob + macros = { + "REPLICAS": str(replicas), + "TASK_NAME": task_name, + } + + # Build head group spec (rank 0) + head_container = copy.deepcopy(base_container) + head_container["name"] = "ray-head" + + # For RayJob, we need to extract and process command/args before removing them + original_command = copy.deepcopy(head_container.get("command", [])) + original_args = copy.deepcopy(head_container.get("args", [])) + + # Apply arg replacements using shared utility + apply_arg_overrides(head_container, task_intent, macros) + + # Update original_command/original_args for entrypoint generation after overrides + original_command = copy.deepcopy(head_container.get("command", [])) + original_args = copy.deepcopy(head_container.get("args", [])) + + # Substitute /results/output paths to point within coldpress base directory + # This ensures training outputs are saved in the coldpress workspace + output_path_substitutions = { + "/results/output": f"/results/{base_dir}/output", + "/results/checkpoints": f"/results/{base_dir}/output", + } + + def substitute_output_paths(args_list): + """Replace output paths in arguments with coldpress base directory paths""" + if not args_list: + return args_list + substituted = [] + for arg in args_list: + new_arg = arg + for old_path, new_path in output_path_substitutions.items(): + if old_path in arg: + new_arg = arg.replace(old_path, new_path) + substituted.append(new_arg) + return substituted + + original_args = substitute_output_paths(original_args) + original_command = substitute_output_paths(original_command) + + if "command" in head_container: + del head_container["command"] + if "args" in head_container: + del head_container["args"] + + # Apply env var additions/replacements using shared utility + head_macros = macros.copy() + head_macros["INDEX"] = "0" + head_macros["RANK"] = "0" + apply_env_overrides(head_container, task_intent, head_macros) + + # Build head pod spec + head_pod_spec = { + "containers": [head_container], + } + + # Copy tolerations + if vk8s_spec.get("tolerations"): + head_pod_spec["tolerations"] = vk8s_spec["tolerations"] + + # Copy volumes and add coldpress- prefix to configMap names + if vk8s_spec.get("volumes"): + head_pod_spec["volumes"] = copy.deepcopy(vk8s_spec["volumes"]) + for volume in head_pod_spec["volumes"]: + if "configMap" in volume: + cm_name = volume["configMap"]["name"] + if not cm_name.startswith("coldpress-"): + volume["configMap"]["name"] = f"coldpress-{cm_name}" + + # Add storage PVC if not already present + if not any(v.get("name") == "results" for v in head_pod_spec.get("volumes", [])): + if "volumes" not in head_pod_spec: + head_pod_spec["volumes"] = [] + head_pod_spec["volumes"].insert( + 0, + {"name": "results", "persistentVolumeClaim": {"claimName": data_pvc_name}}, + ) + + # Build init containers for head + init_containers = [] + + # Validate base_dir contains no shell metacharacters (defense in depth) + for char in DANGEROUS_SHELL_CHARS: + if char in base_dir: + raise ValueError( + f"Invalid character '{char}' in base_dir '{base_dir}'. " + f"This could be a security risk." + ) + + # SAFETY: base_dir validated above, replicas is int - safe for shell command + # Add mkdir init container - create directories for head and all workers + # Also create workspace and output directories for training files and results + mkdir_dirs = ( + f"/results/{base_dir} /results/{base_dir}/workspace /results/{base_dir}/output" + ) + for worker_idx in range(1, replicas): + mkdir_dirs += f" /results/{base_dir}/worker-{worker_idx}" + + mkdir_init = { + "name": "mkdir", + "image": MKDIR_IMAGE, + "command": ["sh", "-c"], + "args": [ + f"mkdir -p {mkdir_dirs} && echo Created directories for head and {replicas - 1} workers" + ], + "volumeMounts": [{"name": "results", "mountPath": "/results"}], + } + init_containers.append(mkdir_init) + + # SAFETY: base_dir validated above - safe for shell command + # Add init container to copy training files to shared storage + # This makes them accessible to Ray job submission + # Use -L to dereference symlinks (ConfigMaps create symlinks) + copy_files_init = { + "name": "copy-files", + "image": base_container["image"], + "command": ["sh", "-c"], + "args": [ + f"cp -rL /workspace/* /results/{base_dir}/workspace/ && echo Copied training files to shared storage" + ], + "volumeMounts": [ + {"name": "results", "mountPath": "/results"}, + {"name": "training-script", "mountPath": "/workspace"}, + ], + } + init_containers.append(copy_files_init) + + # Add discovery init container if specified + should_run_discovery = discovery_template_path and task_name in discovery_task_names + + if should_run_discovery: + # Find storage volume name + storage_volume_name = None + for volume in head_pod_spec.get("volumes", []): + if "persistentVolumeClaim" in volume: + pvc_claim = volume["persistentVolumeClaim"].get("claimName") + if pvc_claim == data_pvc_name: + storage_volume_name = volume.get("name") + break + + if storage_volume_name: + discovery_init = build_discovery_init_container( + discovery_template_path, base_dir, storage_volume_name + ) + if discovery_init: + init_containers.append(discovery_init) + + head_pod_spec["initContainers"] = init_containers + + # Build head group spec + head_group_spec = { + "serviceType": "ClusterIP", + "rayStartParams": { + "dashboard-host": "0.0.0.0", + }, + "template": { + "metadata": { + "annotations": {"sidecar.istio.io/inject": "false"}, + "labels": COLDPRESS_LABELS.copy(), + }, + "spec": head_pod_spec, + }, + } + head_group_spec["template"]["metadata"]["labels"]["coldpress.io/job-id"] = ( + rayjob_name + ) + + # Build worker group specs if replicas > 1 + worker_group_specs = [] + + if replicas > 1: + # For each worker, create a separate worker group with 1 replica + for worker_idx in range(1, replicas): + worker_container = copy.deepcopy(base_container) + worker_container["name"] = "ray-worker" + + # Apply arg replacements using shared utility + worker_macros = macros.copy() + worker_macros["INDEX"] = str(worker_idx) + apply_arg_overrides(worker_container, task_intent, worker_macros) + + # For RayJob, remove command/args - Ray operator controls container startup + if "command" in worker_container: + del worker_container["command"] + if "args" in worker_container: + del worker_container["args"] + + # Apply env var additions/replacements using shared utility + worker_macros = macros.copy() + worker_macros["INDEX"] = str(worker_idx) + worker_macros["RANK"] = str(worker_idx) + apply_env_overrides(worker_container, task_intent, worker_macros) + + # Build worker pod spec + worker_pod_spec = { + "containers": [worker_container], + } + + # Add discovery init container for worker if specified + should_run_discovery_worker = ( + discovery_template_path and task_name in discovery_task_names + ) + + worker_init_containers = [] + + if should_run_discovery_worker: + # Find storage volume name + storage_volume_name_worker = None + if vk8s_spec.get("volumes"): + for volume in vk8s_spec["volumes"]: + if "persistentVolumeClaim" in volume: + pvc_claim = volume["persistentVolumeClaim"].get("claimName") + if pvc_claim == data_pvc_name: + storage_volume_name_worker = volume.get("name") + break + + if storage_volume_name_worker: + # Create worker-specific discovery subdirectory + worker_base_dir = f"{base_dir}/worker-{worker_idx}" + discovery_init_worker = build_discovery_init_container( + discovery_template_path, + worker_base_dir, + storage_volume_name_worker, + ) + if discovery_init_worker: + worker_init_containers.append(discovery_init_worker) + + if worker_init_containers: + worker_pod_spec["initContainers"] = worker_init_containers + + # Copy tolerations + if vk8s_spec.get("tolerations"): + worker_pod_spec["tolerations"] = vk8s_spec["tolerations"] + + # Copy volumes + if vk8s_spec.get("volumes"): + worker_pod_spec["volumes"] = copy.deepcopy(vk8s_spec["volumes"]) + # Add coldpress- prefix to configMap volume names + for volume in worker_pod_spec["volumes"]: + if "configMap" in volume: + cm_name = volume["configMap"]["name"] + if not cm_name.startswith("coldpress-"): + volume["configMap"]["name"] = f"coldpress-{cm_name}" + + # Add storage PVC if not already present + if not any( + v.get("name") == "results" for v in worker_pod_spec.get("volumes", []) + ): + if "volumes" not in worker_pod_spec: + worker_pod_spec["volumes"] = [] + worker_pod_spec["volumes"].insert( + 0, + { + "name": "results", + "persistentVolumeClaim": {"claimName": data_pvc_name}, + }, + ) + + worker_group_spec = { + "groupName": f"worker-group-{worker_idx}", + "replicas": 1, + "minReplicas": 1, + "maxReplicas": 1, + "rayStartParams": {}, + "template": { + "metadata": { + "annotations": {"sidecar.istio.io/inject": "false"}, + "labels": COLDPRESS_LABELS.copy(), + }, + "spec": worker_pod_spec, + }, + } + worker_group_spec["template"]["metadata"]["labels"][ + "coldpress.io/job-id" + ] = rayjob_name + + worker_group_specs.append(worker_group_spec) + + # Build RayCluster spec + ray_cluster_spec = { + "headGroupSpec": head_group_spec, + } + + if worker_group_specs: + ray_cluster_spec["workerGroupSpecs"] = worker_group_specs + + # Build RayJob labels + rayjob_labels = COLDPRESS_LABELS.copy() + rayjob_labels.update( + { + "coldpress.io/job-id": rayjob_name, + "kueue.x-k8s.io/queue-name": get_kueue_queue_label(namespace), + } + ) + + # Build entrypoint from original command/args + # For RayJob, use shared storage location where files were copied + shared_workspace = f"/results/{base_dir}/workspace" + + # Build env var substitution map for entrypoint script + # RayJob entrypoint runs before Ray connection, so runtimeEnv vars aren't available + # We need to substitute the values directly into the bash script + # Use all env vars from head container (includes both job-spec defaults and intent overrides) + env_substitutions = {} + env_list = head_container.get("env", []) + for env_item in env_list: + env_key = env_item.get("name") + env_value = env_item.get("value", "") + if env_key: + env_substitutions[f"${env_key}"] = str(env_value) + + entrypoint_cmd = "" + + if original_command == ["python"]: + # Common pattern: command: ["python"], args: ["script.py", "--arg1", ...] + if original_args: + entrypoint_cmd = "python " + " ".join(original_args) + elif original_command in [["bash", "-c"], ["sh", "-c"]]: + # Shell pattern: command: ["bash", "-c"], args: ["script"] + # Need to preserve bash -c wrapper for RayJob entrypoint + # Prepend cd to shared workspace directory + # Also substitute env var references with actual values + if original_args: + shell = "bash" if original_command == ["bash", "-c"] else "sh" + script = original_args[0] + + # Substitute env var references in the script + for var_ref, value in env_substitutions.items(): + script = script.replace(var_ref, value) + + script_with_cd = f"cd {shared_workspace} && {script}" + entrypoint_cmd = f"{shell} -c '{script_with_cd}'" + else: + # General case: combine command + args + if original_command: + entrypoint_cmd = " ".join(original_command) + if original_args: + args_str = " ".join(original_args) + if entrypoint_cmd: + entrypoint_cmd = f"{entrypoint_cmd} {args_str}" + else: + entrypoint_cmd = args_str + + # Prepend cd to shared workspace directory for non-shell commands + # Wrap in bash -c since RayJob entrypoint doesn't support shell operators + if original_command not in [["bash", "-c"], ["sh", "-c"]]: + entrypoint_cmd = f"cd {shared_workspace} && {entrypoint_cmd}" + entrypoint_cmd = f"bash -c '{entrypoint_cmd}'" + + # Build runtime env with all env vars from head container + # (includes both job-spec defaults and intent overrides) + runtime_env = {} + env_list = head_container.get("env", []) + if env_list: + env_vars = {} + for env_item in env_list: + env_key = env_item.get("name") + env_value = env_item.get("value", "") + if env_key: + env_vars[env_key] = str(env_value) + + if env_vars: + runtime_env["env_vars"] = env_vars + + # Build RayJob manifest + rayjob = { + "apiVersion": "ray.io/v1", + "kind": "RayJob", + "metadata": { + "name": rayjob_name, + "namespace": namespace, + "labels": rayjob_labels, + "annotations": { + "coldpress.io/base-dir": base_dir, + "coldpress.io/storage-pvc": data_pvc_name, + }, + }, + "spec": { + "suspend": True, # Start suspended for Kueue + "entrypoint": entrypoint_cmd, + "rayClusterSpec": ray_cluster_spec, + "shutdownAfterJobFinishes": True, + "ttlSecondsAfterFinished": 3600, # Clean up after 1 hour + }, + } + + # Add runtime env if we have env vars to pass + if runtime_env: + rayjob["spec"]["runtimeEnvYAML"] = yaml.safe_dump(runtime_env) + + return rayjob, base_dir + + +def kuberay_to_yaml(manifest): + """Convert KubeRay manifest dict to YAML string. + + Args: + manifest: KubeRay manifest dict + + Returns: + str: YAML string + """ + return yaml.safe_dump(manifest, default_flow_style=False, sort_keys=False) diff --git a/coldpress/script_gen.py b/coldpress/script_gen.py index fae459b..a76c558 100644 --- a/coldpress/script_gen.py +++ b/coldpress/script_gen.py @@ -1,34 +1,198 @@ -# Generated by: Claude Sonnet 4.5 +# Assisted by: Claude Sonnet 4.5 """Bash script generation for Coldpress jobs.""" +import os from datetime import datetime +from .constants import ( + DEFAULT_JOB_TIMEOUT, + DEFAULT_SLEEP_DURATION, + DEFAULT_SLEEP_INFINITY, + EXPLORER_IMAGE, + COPIER_IMAGE, + MKDIR_IMAGE, + MANIFEST_CONFIG, +) + + +def sanitize_identifier(identifier: str, identifier_type: str = "identifier") -> str: + """ + Sanitize an identifier (job name, namespace) for safe use in shell commands. + + Args: + identifier: Identifier to sanitize + identifier_type: Type of identifier for error messages (e.g., "job name", "namespace") + + Returns: + Sanitized identifier + + Raises: + ValueError: If identifier contains dangerous characters + """ + if not identifier or not isinstance(identifier, str): + raise ValueError(f"Invalid {identifier_type}: {identifier}") + + # Check for dangerous shell metacharacters + dangerous_chars = [ + "`", + "$", + "|", + ";", + "&", + ">", + "<", + "\n", + "\r", + " ", + "'", + '"', + "/", + "\\", + ] + for char in dangerous_chars: + if char in identifier: + raise ValueError( + f"Invalid character '{char}' in {identifier_type} '{identifier}'. " + f"This could be a security risk." + ) + + # Ensure it's not empty + if not identifier or identifier in [".", ".."]: + raise ValueError(f"Invalid {identifier_type}: '{identifier}'") + + return identifier + + +def sanitize_filename(filename: str) -> str: + """ + Sanitize a filename for safe use in shell commands. + + Args: + filename: Filename to sanitize + + Returns: + Sanitized filename + + Raises: + ValueError: If filename contains dangerous characters or path traversal + """ + if not filename or not isinstance(filename, str): + raise ValueError(f"Invalid filename: {filename}") + + # Get basename to prevent path traversal + basename = os.path.basename(filename) + if basename != filename: + raise ValueError( + f"Filename contains path separators: '{filename}'. " + "Only basenames are allowed." + ) + + # Check for dangerous shell metacharacters + dangerous_chars = ["`", "$", "|", ";", "&", ">", "<", "\n", "\r", " ", "'", '"'] + for char in dangerous_chars: + if char in basename: + raise ValueError( + f"Invalid character '{char}' in filename '{basename}'. " + f"This could be a security risk." + ) + + # Ensure it's not empty after basename + if not basename or basename in [".", ".."]: + raise ValueError(f"Invalid filename: '{basename}'") + + return basename def generate_run_script( - job_name, namespace, timeout="1h", configmap_name=None, configmap_files=None + job_name, + namespace, + timeout=None, + configmap_name=None, + configmap_files=None, + manifest_type="jobset", ): """ Generate run.sh script to apply and monitor the job. Args: - job_name: Name of the JobSet + job_name: Name of the job namespace: Kubernetes namespace - timeout: Timeout for job completion (default: 1h) + timeout: Timeout for job completion (default: from DEFAULT_JOB_TIMEOUT) configmap_name: Optional ConfigMap name configmap_files: Optional list of files to include in ConfigMap + manifest_type: Type of manifest (jobset, pytorchjob, tfjob, mpijob, inferenceservice) Returns: str: Bash script content """ + # Validate inputs to prevent shell injection + sanitize_identifier(job_name, "job name") + sanitize_identifier(namespace, "namespace") + + if timeout is None: + timeout = DEFAULT_JOB_TIMEOUT + configmap_apply = "" if configmap_name and configmap_files: - from_files = " ".join([f"--from-file={f}" for f in configmap_files]) + # Sanitize filenames to prevent command injection + sanitized_files = [] + for f in configmap_files: + try: + safe_filename = sanitize_filename(f) + sanitized_files.append(safe_filename) + except ValueError as e: + raise ValueError(f"Invalid ConfigMap filename: {e}") from e + + from_files = " ".join([f"--from-file={f}" for f in sanitized_files]) + # Add coldpress- prefix if not already present + if not configmap_name.startswith("coldpress-"): + actual_configmap_name = f"coldpress-{configmap_name}" + else: + actual_configmap_name = configmap_name configmap_apply = f""" # Create ConfigMap from files echo "Creating ConfigMap from files..." -oc create configmap {configmap_name} -n {namespace} {from_files} --dry-run=client -o yaml | oc apply -f - +oc create configmap {actual_configmap_name} -n {namespace} {from_files} --dry-run=client -o yaml | oc apply -f - """ + # Get config or use defaults for unknown types + config = MANIFEST_CONFIG.get( + manifest_type, + { + "file": "manifest.yaml", + "type": manifest_type, + "apply_msg": manifest_type, + }, + ) + + manifest_file = config["file"] + resource_type = config["type"] + resource_name = f"coldpress-{job_name}" + apply_cmd = f'echo "Applying {config["apply_msg"]}..."\noc apply -f {manifest_file}' + + # Generate services apply command if applicable + services_apply = ( + """ +# Apply services if they exist +if [ -f services.yaml ]; then + echo "Applying Services..." + oc apply -f services.yaml +fi""" + if config.get("has_services", False) + else "" + ) + + # Determine wait condition based on resource type + if manifest_type == "inferenceservice": + wait_message = ( + f"Waiting for InferenceService to be ready (timeout: {timeout})..." + ) + wait_cmd = f"oc wait --for=condition=Ready $RESOURCE_TYPE/$JOB_NAME -n $NAMESPACE --timeout={timeout} 2>/dev/null" + pod_selector = f"serving.kserve.io/inferenceservice={job_name}" + else: + wait_message = f"Waiting for job to complete (timeout: {timeout})..." + wait_cmd = f"oc wait --for=condition=complete $RESOURCE_TYPE/$JOB_NAME -n $NAMESPACE --timeout={timeout} 2>/dev/null" + pod_selector = f"coldpress.io/job-id={job_name}" + script = f'''#!/bin/bash # Generated by coldpress on {datetime.now().strftime("%Y-%m-%d %H:%M:%S")} @@ -38,55 +202,92 @@ def generate_run_script( cd "$(dirname "$0")" NAMESPACE="{namespace}" -JOBSET_NAME="{job_name}" +JOB_NAME="{resource_name}" +RESOURCE_TYPE="{resource_type}" {configmap_apply} -echo "Applying JobSet..." -oc apply -f jobset.yaml +{apply_cmd} +{services_apply} -# Apply services if they exist -if [ -f services.yaml ]; then - echo "Applying Services..." - oc apply -f services.yaml -fi - -echo "Waiting for job to complete (timeout: {timeout})..." -oc wait --for=condition=complete jobset/$JOBSET_NAME -n $NAMESPACE --timeout={timeout} || {{ - echo "Job did not complete within {timeout}" +echo "{wait_message}" +{wait_cmd} || {{ + echo "Job may still be running or condition not supported for $RESOURCE_TYPE" echo "Current status:" - oc get jobset/$JOBSET_NAME -n $NAMESPACE - oc get jobs -n $NAMESPACE -l coldpress/gid=$JOBSET_NAME - oc get pods -n $NAMESPACE -l coldpress/gid=$JOBSET_NAME - exit 1 + oc get $RESOURCE_TYPE/$JOB_NAME -n $NAMESPACE + oc get pods -n $NAMESPACE -l training.kubeflow.org/job-name=$JOB_NAME 2>/dev/null || \ + oc get pods -n $NAMESPACE -l {pod_selector} 2>/dev/null || \ + oc get pods -n $NAMESPACE | grep {job_name} }} -echo "Job completed successfully!" +echo "Job applied successfully!" echo "" -echo "Results are stored in PVC. To extract results, run:" -echo " oc exec -n $NAMESPACE deployment/ -- tar czf - /mnt/coldpress-data/coldpress_results/$JOBSET_NAME | tar xzf -" +echo "To monitor: ./monitor.sh" +echo "To check logs: ./logs.sh" +echo "To cleanup: ./cleanup.sh" ''' return script -def generate_cleanup_script(job_name, namespace, configmap_name=None): +def generate_cleanup_script( + job_name, namespace, configmap_name=None, manifest_type="jobset" +): """ Generate cleanup.sh script to delete job resources. Args: - job_name: Name of the JobSet + job_name: Name of the job (without coldpress- prefix) namespace: Kubernetes namespace configmap_name: Optional ConfigMap name to delete + manifest_type: Type of manifest (jobset, pytorchjob, tfjob, mpijob, inferenceservice) Returns: str: Bash script content """ + # Validate inputs to prevent shell injection + sanitize_identifier(job_name, "job name") + sanitize_identifier(namespace, "namespace") + configmap_delete = "" if configmap_name: + # Add coldpress- prefix if not already present + if not configmap_name.startswith("coldpress-"): + actual_configmap_name = f"coldpress-{configmap_name}" + else: + actual_configmap_name = configmap_name configmap_delete = f""" # Delete ConfigMap echo "Deleting ConfigMap..." -oc delete configmap/{configmap_name} -n $NAMESPACE --ignore-not-found=true +oc delete configmap/{actual_configmap_name} -n $NAMESPACE --ignore-not-found=true """ + # Generate manifest-specific commands + if manifest_type == "jobset": + resource_type = "jobset" + resource_name = f"coldpress-{job_name}" + delete_cmd = f"oc delete {resource_type}/{resource_name} -n $NAMESPACE --ignore-not-found=true" + services_delete = f"oc delete services -n $NAMESPACE -l coldpress/gid={resource_name} --ignore-not-found=true" + elif manifest_type in ["pytorchjob", "tfjob", "mpijob"]: + resource_name = f"coldpress-{job_name}" + delete_cmd = f"oc delete {manifest_type}/{resource_name} -n $NAMESPACE --ignore-not-found=true" + services_delete = "" # Kubeflow jobs don't create separate services + elif manifest_type == "inferenceservice": + resource_name = f"coldpress-{job_name}" + delete_cmd = f"oc delete {manifest_type}/{resource_name} -n $NAMESPACE --ignore-not-found=true" + services_delete = "" + else: + resource_name = f"coldpress-{job_name}" + delete_cmd = f"oc delete {manifest_type}/{resource_name} -n $NAMESPACE --ignore-not-found=true" + services_delete = "" + + services_section = ( + f""" +# Delete Services +echo "Deleting Services..." +{services_delete} +""" + if services_delete + else "" + ) + script = f'''#!/bin/bash # Generated by coldpress on {datetime.now().strftime("%Y-%m-%d %H:%M:%S")} @@ -96,51 +297,80 @@ def generate_cleanup_script(job_name, namespace, configmap_name=None): cd "$(dirname "$0")" NAMESPACE="{namespace}" -JOBSET_NAME="{job_name}" +JOB_NAME="{resource_name}" +RESOURCE_TYPE="{manifest_type}" -echo "Cleaning up resources for job: $JOBSET_NAME" +echo "Cleaning up resources for job: $JOB_NAME" -# Delete JobSet (this cascades to Jobs and Pods) -echo "Deleting JobSet..." -oc delete jobset/$JOBSET_NAME -n $NAMESPACE --ignore-not-found=true - -# Delete Services -echo "Deleting Services..." -oc delete services -n $NAMESPACE -l coldpress/gid=$JOBSET_NAME --ignore-not-found=true -{configmap_delete} +# Delete main resource (this cascades to Jobs and Pods) +echo "Deleting $RESOURCE_TYPE..." +{delete_cmd} +{services_section}{configmap_delete} # Delete helper pods (log-saver, coldpress-explorer) echo "Deleting helper pods..." -oc delete pods -n $NAMESPACE -l app=coldpress-explorer,coldpress/gid=$JOBSET_NAME --ignore-not-found=true -oc get pods -n $NAMESPACE --no-headers | grep -E "log-saver-${{JOBSET_NAME}}-|coldpress-explorer-${{JOBSET_NAME}}-" | awk '{{print $1}}' | xargs -r oc delete pod -n $NAMESPACE --ignore-not-found=true +oc delete pods -n $NAMESPACE -l app=coldpress-explorer,coldpress.io/job-id=$JOB_NAME --ignore-not-found=true 2>/dev/null || true +oc get pods -n $NAMESPACE --no-headers 2>/dev/null | grep -E "log-saver-${{JOB_NAME}}-|coldpress-explorer-${{JOB_NAME}}-|coldpress-copier-${{JOB_NAME}}-" | awk '{{print $1}}' | xargs -r oc delete pod -n $NAMESPACE --ignore-not-found=true 2>/dev/null || true echo "Cleanup complete!" ''' return script -def generate_monitor_script(job_name, namespace): +def generate_monitor_script(job_name, namespace, manifest_type="jobset"): """ Generate monitor.sh script to watch job progress. Args: - job_name: Name of the JobSet + job_name: Name of the job namespace: Kubernetes namespace + manifest_type: Type of manifest (jobset, pytorchjob, etc.) Returns: str: Bash script content """ + # Validate inputs to prevent shell injection + sanitize_identifier(job_name, "job name") + sanitize_identifier(namespace, "namespace") + + # Generate manifest-specific commands + if manifest_type == "jobset": + resource_type = "jobset" + resource_name = f"coldpress-{job_name}" + label_selector = f"coldpress.io/job-id={resource_name}" + elif manifest_type == "pytorchjob": + resource_type = "pytorchjob" + resource_name = f"coldpress-{job_name}" + label_selector = f"training.kubeflow.org/job-name={resource_name}" + elif manifest_type == "tfjob": + resource_type = "tfjob" + resource_name = f"coldpress-{job_name}" + label_selector = f"training.kubeflow.org/job-name={resource_name}" + elif manifest_type == "mpijob": + resource_type = "mpijob" + resource_name = f"coldpress-{job_name}" + label_selector = f"training.kubeflow.org/job-name={resource_name}" + elif manifest_type == "inferenceservice": + resource_type = "inferenceservice" + resource_name = f"coldpress-{job_name}" + label_selector = f"serving.kserve.io/inferenceservice={resource_name}" + else: + resource_type = manifest_type + resource_name = f"coldpress-{job_name}" + label_selector = f"coldpress.io/job-id={resource_name}" + script = f'''#!/bin/bash # Generated by coldpress on {datetime.now().strftime("%Y-%m-%d %H:%M:%S")} NAMESPACE="{namespace}" -JOBSET_NAME="{job_name}" +JOB_NAME="{resource_name}" +RESOURCE_TYPE="{resource_type}" -echo "Monitoring job: $JOBSET_NAME in namespace: $NAMESPACE" +echo "Monitoring $RESOURCE_TYPE: $JOB_NAME in namespace: $NAMESPACE" echo "Press Ctrl+C to exit" echo "" -watch -n 2 "oc get jobset,job,pod -n $NAMESPACE -l coldpress/gid=$JOBSET_NAME" +watch -n 2 "oc get $RESOURCE_TYPE,job,pod -n $NAMESPACE -l {label_selector}" ''' return script @@ -150,7 +380,7 @@ def generate_logs_script(job_name, namespace, storage_pvc, base_dir): Generate logs.sh script to save and view logs from all pods. Args: - job_name: Name of the JobSet + job_name: Name of the job (without coldpress- prefix) namespace: Kubernetes namespace storage_pvc: PVC name for results storage base_dir: Base directory path in PVC @@ -158,28 +388,38 @@ def generate_logs_script(job_name, namespace, storage_pvc, base_dir): Returns: str: Bash script content """ + # Validate inputs to prevent shell injection + sanitize_identifier(job_name, "job name") + sanitize_identifier(namespace, "namespace") + sanitize_identifier(storage_pvc, "storage PVC") + + full_job_name = f"coldpress-{job_name}" + script = f'''#!/bin/bash # Generated by coldpress on {datetime.now().strftime("%Y-%m-%d %H:%M:%S")} NAMESPACE="{namespace}" -JOBSET_NAME="{job_name}" +JOB_NAME="{full_job_name}" STORAGE_PVC="{storage_pvc}" BASE_DIR="{base_dir}" LOGS_DIR="$BASE_DIR/logs" -echo "Fetching and saving logs for job: $JOBSET_NAME" +echo "Fetching and saving logs for job: $JOB_NAME" echo "" -# Get all pods for this job -PODS=$(oc get pods -n $NAMESPACE -l coldpress/gid=$JOBSET_NAME -o jsonpath='{{.items[*].metadata.name}}') +# Get all pods for this job (try both label formats for compatibility) +PODS=$(oc get pods -n $NAMESPACE -l coldpress.io/job-id=$JOB_NAME -o jsonpath='{{.items[*].metadata.name}}' 2>/dev/null) +if [ -z "$PODS" ]; then + PODS=$(oc get pods -n $NAMESPACE -l coldpress/gid=$JOB_NAME -o jsonpath='{{.items[*].metadata.name}}' 2>/dev/null) +fi if [ -z "$PODS" ]; then - echo "No pods found for job $JOBSET_NAME" + echo "No pods found for job $JOB_NAME" exit 1 fi # Create helper pod to save logs to PVC -POD_NAME="log-saver-${{JOBSET_NAME}}-$(date +%s)" +POD_NAME="log-saver-${{JOB_NAME}}-$(date +%s)" # Cleanup function cleanup() {{ @@ -204,8 +444,8 @@ def generate_logs_script(job_name, namespace, storage_pvc, base_dir): restartPolicy: Never containers: - name: saver - image: registry.access.redhat.com/ubi9/ubi-minimal:latest - command: ["sleep", "300"] + image: {MKDIR_IMAGE} + command: ["sleep", "{DEFAULT_SLEEP_DURATION}"] volumeMounts: - name: data mountPath: /data @@ -256,7 +496,7 @@ def generate_explore_script(job_name, namespace, storage_pvc, base_dir): Automatically cleans up the pod when the shell session exits. Args: - job_name: Name of the JobSet + job_name: Name of the job (without coldpress- prefix) namespace: Kubernetes namespace storage_pvc: Name of the storage PVC base_dir: Base directory path in PVC @@ -264,16 +504,23 @@ def generate_explore_script(job_name, namespace, storage_pvc, base_dir): Returns: str: Bash script content """ + # Validate inputs to prevent shell injection + sanitize_identifier(job_name, "job name") + sanitize_identifier(namespace, "namespace") + sanitize_identifier(storage_pvc, "storage PVC") + + full_job_name = f"coldpress-{job_name}" + script = f'''#!/bin/bash # Generated by coldpress on {datetime.now().strftime("%Y-%m-%d %H:%M:%S")} set -e NAMESPACE="{namespace}" -JOBSET_NAME="{job_name}" +JOB_NAME="{full_job_name}" STORAGE_PVC="{storage_pvc}" BASE_DIR="{base_dir}" -POD_NAME="coldpress-explorer-${{JOBSET_NAME}}-$(date +%s)" +POD_NAME="coldpress-explorer-${{JOB_NAME}}-$(date +%s)" echo "Starting interactive explorer pod..." echo "PVC: $STORAGE_PVC" @@ -289,13 +536,13 @@ def generate_explore_script(job_name, namespace, storage_pvc, base_dir): namespace: $NAMESPACE labels: app: coldpress-explorer - coldpress/gid: $JOBSET_NAME + coldpress.io/job-id: $JOB_NAME spec: restartPolicy: Never containers: - name: explorer - image: registry.access.redhat.com/ubi9/ubi-minimal:latest - command: ["sleep", "infinity"] + image: {EXPLORER_IMAGE} + command: ["sleep", "{DEFAULT_SLEEP_INFINITY}"] volumeMounts: - name: data mountPath: /data @@ -350,7 +597,7 @@ def generate_copy_script(job_name, namespace, storage_pvc, base_dir): Automatically cleans up the pod when the copy is complete. Args: - job_name: Name of the JobSet + job_name: Name of the job (without coldpress- prefix) namespace: Kubernetes namespace storage_pvc: Name of the storage PVC base_dir: Base directory path in PVC @@ -358,16 +605,23 @@ def generate_copy_script(job_name, namespace, storage_pvc, base_dir): Returns: str: Bash script content """ + # Validate inputs to prevent shell injection + sanitize_identifier(job_name, "job name") + sanitize_identifier(namespace, "namespace") + sanitize_identifier(storage_pvc, "storage PVC") + + full_job_name = f"coldpress-{job_name}" + script = f'''#!/bin/bash # Generated by coldpress on {datetime.now().strftime("%Y-%m-%d %H:%M:%S")} set -e NAMESPACE="{namespace}" -JOBSET_NAME="{job_name}" +JOB_NAME="{full_job_name}" STORAGE_PVC="{storage_pvc}" BASE_DIR="{base_dir}" -POD_NAME="coldpress-copier-${{JOBSET_NAME}}-$(date +%s)" +POD_NAME="coldpress-copier-${{JOB_NAME}}-$(date +%s)" # Determine destination directory DEST_DIR="${{1:-$(dirname "$0")/results}}" @@ -388,13 +642,13 @@ def generate_copy_script(job_name, namespace, storage_pvc, base_dir): namespace: $NAMESPACE labels: app: coldpress-copier - coldpress/gid: $JOBSET_NAME + coldpress.io/job-id: $JOB_NAME spec: restartPolicy: Never containers: - name: copier - image: registry.access.redhat.com/ubi9/ubi:latest - command: ["sleep", "300"] + image: {COPIER_IMAGE} + command: ["sleep", "{DEFAULT_SLEEP_DURATION}"] volumeMounts: - name: data mountPath: /data @@ -442,18 +696,20 @@ def write_scripts( base_dir=None, configmap_name=None, configmap_files=None, + manifest_type="jobset", ): """ Write all bash scripts to output directory. Args: output_dir: Directory to write scripts to - job_name: Name of the JobSet + job_name: Name of the job namespace: Kubernetes namespace storage_pvc: Name of the storage PVC (optional) base_dir: Base directory path in PVC (optional) configmap_name: Name of ConfigMap to apply/delete (optional) configmap_files: List of files for ConfigMap (optional) + manifest_type: Type of manifest (jobset, pytorchjob, etc.) """ import os @@ -463,11 +719,17 @@ def write_scripts( namespace, configmap_name=configmap_name, configmap_files=configmap_files, + manifest_type=manifest_type, ), "cleanup.sh": generate_cleanup_script( - job_name, namespace, configmap_name=configmap_name + job_name, + namespace, + configmap_name=configmap_name, + manifest_type=manifest_type, + ), + "monitor.sh": generate_monitor_script( + job_name, namespace, manifest_type=manifest_type ), - "monitor.sh": generate_monitor_script(job_name, namespace), } # Add logs, explore, and copy scripts if storage info is available diff --git a/coldpress/utils.py b/coldpress/utils.py new file mode 100644 index 0000000..2d80bf6 --- /dev/null +++ b/coldpress/utils.py @@ -0,0 +1,553 @@ +# Assisted by: Claude Sonnet 4.5 +"""Shared utility functions for Coldpress generators. + +This module contains common functions used across JobSet, Kubeflow, and KubeRay generators +to reduce code duplication and ensure consistent behavior. +""" + +import os +import re +import copy +import sys +import yaml + +# Security: Dangerous shell metacharacters that could enable injection +DANGEROUS_SHELL_CHARS = ["`", "$", "|", ";", "&", ">", "<", "\n", "\r"] + +# Discovery container mount path (Kubernetes volume mount point) +DISCOVERY_MOUNT_PATH = "/tmp/result" + + +def sanitize_value(value: str, context: str = "") -> str: + """ + Sanitize a value for safe use in Kubernetes manifests. + + Args: + value: Value to sanitize + context: Context for error messages + + Returns: + Sanitized value + + Raises: + ValueError: If value contains dangerous characters + """ + if not isinstance(value, str): + return str(value) + + # Check for shell metacharacters that could be dangerous in args + for char in DANGEROUS_SHELL_CHARS: + if char in value: + raise ValueError( + f"Invalid character '{char}' in value '{value}'{context}. " + f"This could be a security risk." + ) + + return value + + +def substitute_macros(value: str, macros: dict[str, str], context: str = "") -> str: + """ + Substitute ${MACRO} patterns in a string with sanitization. + + Args: + value: String potentially containing ${MACRO} patterns + macros: Dict of macro names to values + context: Context for error messages (e.g., " for arg 'master_addr'") + + Returns: + String with macros substituted + + Raises: + ValueError: If macro value contains invalid characters + """ + if not isinstance(value, str): + return value + + pattern = r"\$\{([^}]+)\}" + + def replacer(match): + macro_name = match.group(1) + if macro_name not in macros: + # Leave unresolved macros as-is + return match.group(0) + + macro_value = str(macros[macro_name]) + + # Sanitize the macro value + try: + sanitized = sanitize_value(macro_value, f"{context} (macro: {macro_name})") + except ValueError as e: + raise ValueError(f"Macro substitution failed: {e}") from e + + return sanitized + + result = re.sub(pattern, replacer, value) + + # Final sanitization of the complete result + return sanitize_value(result, context) + + +def parse_arg(arg: str) -> tuple: + """ + Parse a command-line argument into (flag, value) tuple. + + Returns: + (flag, value): For --flag=value or --flag + (None, value): For positional args + """ + arg_str = str(arg) + if arg_str.startswith("--") or arg_str.startswith("-"): + if "=" in arg_str: + flag, value = arg_str.split("=", 1) + return (flag, value) + else: + return (arg_str, None) + else: + return (None, arg_str) + + +def replace_or_add_arg( + args: list, + key: str, + value: str, + insert_after: str = None, + insert_before: str = None, +) -> list: + """ + Replace existing arg or add new arg with optional positioning. + + Args: + args: List of command-line arguments + key: Argument flag (without -- prefix) + value: Argument value + insert_after: Flag to insert after (optional, for new args) + insert_before: Flag to insert before (optional, for new args) + + Returns: + Modified args list + """ + result = list(args) + flag = f"--{key}" + skip_next = False + found = False + + # Phase 1: Try to replace existing arg + for i in range(len(result)): + if skip_next: + skip_next = False + continue + + arg_flag, arg_value = parse_arg(result[i]) + + if arg_flag == flag: + # Found the arg + found = True + if arg_value is not None: + # Format: --flag=value + result[i] = f"{flag}={value}" + else: + # Format: --flag value (two args) + if i + 1 < len(result): + next_flag, next_value = parse_arg(result[i + 1]) + if next_flag is None: # Next is a value + result[i + 1] = value + skip_next = True + else: + # Flag without value, convert to --flag=value + result[i] = f"{flag}={value}" + else: + result[i] = f"{flag}={value}" + break + + # Phase 2: If not found, insert at specified position + if not found: + new_arg = f"{flag}={value}" + + if insert_after: + # Find anchor and insert after it + anchor_flag = ( + insert_after if insert_after.startswith("--") else f"--{insert_after}" + ) + for i, arg in enumerate(result): + arg_flag, _ = parse_arg(arg) + if arg_flag == anchor_flag: + result.insert(i + 1, new_arg) + return result + # Anchor not found, append to end + result.append(new_arg) + + elif insert_before: + # Find anchor and insert before it + anchor_flag = ( + insert_before + if insert_before.startswith("--") + else f"--{insert_before}" + ) + for i, arg in enumerate(result): + arg_flag, _ = parse_arg(arg) + if arg_flag == anchor_flag: + result.insert(i, new_arg) + return result + # Anchor not found, append to end + result.append(new_arg) + + else: + # No position specified, append to end + result.append(new_arg) + + return result + + +def apply_arg_overrides(container: dict, task_intent, macros: dict[str, str]) -> None: + """ + Apply argument overrides from task intent to container spec. + + Args: + container: Container spec dict (modified in-place) + task_intent: Task intent with args field + macros: Macro substitution dict + + Modifies container in-place to update args or command field. + """ + if not task_intent.args: + return + + # Determine which field contains the arguments + if container.get("args"): + args_field = "args" + args = ( + container["args"] + if isinstance(container["args"], list) + else [container["args"]] + ) + elif container.get("command"): + # If no args field but command exists, apply overrides to command + args_field = "command" + args = ( + container["command"] + if isinstance(container["command"], list) + else [container["command"]] + ) + else: + args_field = None + args = None + + if args: + for arg_key, arg_override in task_intent.args.items(): + # Handle both simple string and ArgOverride object + if isinstance(arg_override, str): + # Simple string value + value = arg_override + insert_after = None + insert_before = None + else: + # ArgOverride object (pydantic model) with insertion position + value = arg_override.value + insert_after = arg_override.insert_after + insert_before = arg_override.insert_before + + # Substitute macros with validation + context = f" for arg '--{arg_key}'" + substituted_value = substitute_macros(value, macros, context) + args = replace_or_add_arg( + args, + arg_key, + substituted_value, + insert_after=insert_after, + insert_before=insert_before, + ) + + container[args_field] = args + + +def apply_env_overrides(container: dict, task_intent, macros: dict[str, str]) -> None: + """ + Apply environment variable overrides from task intent to container spec. + + Args: + container: Container spec dict (modified in-place) + task_intent: Task intent with env field + macros: Macro substitution dict + + Modifies container in-place to update env field. + """ + if not task_intent.env: + return + + # Get existing env vars + env_list = container.get("env", []) + if not isinstance(env_list, list): + env_list = [] + + for env_key, env_value in task_intent.env.items(): + # Substitute macros with validation + context = f" for env var '{env_key}'" + substituted_value = substitute_macros(str(env_value), macros, context) + + # Find and replace or append + found = False + for env_item in env_list: + if env_item.get("name") == env_key: + env_item["value"] = substituted_value + found = True + break + + if not found: + env_list.append({"name": env_key, "value": substituted_value}) + + container["env"] = env_list + + +def get_storage_pvc_name(project_config: dict, namespace: str) -> str: + """ + Get storage PVC name from project config with fallback. + + Uses the fallback chain: project config > default PVC name. + + Args: + project_config: Project configuration dict + namespace: Target namespace + + Returns: + PVC name for storage + """ + from .constants import get_pvc_name + + # Storage is validated by Pydantic - if present, 'results' is required + storage = project_config.get("storage") + if storage: + return storage["results"] + else: + return get_pvc_name(namespace) + + +def parse_discovery_config(intent_config): + """ + Extract discovery template path and task names from intent config. + + Args: + intent_config: Validated IntentConfig object + + Returns: + tuple: (discovery_template_path, discovery_task_names) + - discovery_template_path: Full path to discovery template YAML, or None + - discovery_task_names: Set of task names that should run discovery + """ + from .constants import DEFAULT_DISCOVERY_DIR + + discovery_template_path = None + discovery_task_names = set() + + if intent_config.discovery: + discovery_config = intent_config.discovery + template_name = ( + discovery_config.template + if hasattr(discovery_config, "template") + else discovery_config + ) + + # Resolve template path + discovery_dir = os.getenv("COLDPRESS_DISCOVERY_DIR", DEFAULT_DISCOVERY_DIR) + discovery_template_path = os.path.join(discovery_dir, f"{template_name}.yaml") + + # Determine which tasks should run discovery + discovery_tasks = ( + discovery_config.tasks if hasattr(discovery_config, "tasks") else "all" + ) + + if discovery_tasks == "all": + discovery_task_names = {task.name for task in intent_config.tasks} + elif isinstance(discovery_tasks, list): + # discovery_tasks is a list of task names + discovery_task_names = set(discovery_tasks) + + return discovery_template_path, discovery_task_names + + +def extract_configmap_name(manifest: dict, manifest_type: str) -> str | None: + """ + Extract ConfigMap name from a manifest. + + Args: + manifest: The manifest dict (JobSet, PyTorchJob, RayJob, etc.) + manifest_type: Type of manifest ("jobset", "pytorchjob", "rayjob", etc.) + + Returns: + ConfigMap name if found, None otherwise + """ + if manifest_type == "jobset": + # Search for configMap volumes in JobSet manifest + for replicated_job in manifest.get("spec", {}).get("replicatedJobs", []): + job_template = ( + replicated_job.get("template", {}).get("spec", {}).get("template", {}) + ) + volumes = job_template.get("spec", {}).get("volumes", []) + for volume in volumes: + if "configMap" in volume: + return volume["configMap"]["name"] + + elif manifest_type == "pytorchjob": + # Search for configMap volumes in PyTorchJob manifest + replica_specs = manifest.get("spec", {}).get("pytorchReplicaSpecs", {}) + for replica_spec in replica_specs.values(): + volumes = ( + replica_spec.get("template", {}).get("spec", {}).get("volumes", []) + ) + for volume in volumes: + if "configMap" in volume: + return volume["configMap"]["name"] + + elif manifest_type == "rayjob": + # Search for configMap volumes in RayJob manifest (head group) + head_spec = ( + manifest.get("spec", {}).get("rayClusterSpec", {}).get("headGroupSpec", {}) + ) + volumes = head_spec.get("template", {}).get("spec", {}).get("volumes", []) + for volume in volumes: + if "configMap" in volume: + return volume["configMap"]["name"] + + return None + + +def build_discovery_init_container( + template_path: str, + base_dir: str, + storage_volume_name: str, + task_id: int | None = None, + per_pod_directory: bool = False, +): + """ + Build discovery init container from template. + + Args: + template_path: Path to discovery template YAML + base_dir: Base directory path + storage_volume_name: Name of the storage volume to mount + task_id: Task ID for result path (optional, for JobSet use) + per_pod_directory: If True, create pod-specific subdirectories using downward API + + Returns: + dict: Init container spec, or None if template not found + """ + try: + # Read discovery template + with open(template_path, "r") as f: + template = yaml.safe_load(f) + + # Extract Pod spec + pod_spec = template.get("spec", {}) + containers = pod_spec.get("containers", []) + + if not containers: + return None + + # Get template name from filename and sanitize for shell use + template_name = ( + os.path.basename(template_path).replace(".yaml", "").replace(".yml", "") + ) + # Sanitize template_name to prevent command injection + sanitized_template_name = sanitize_value( + template_name, " for discovery template" + ) + + # Copy container and modify for init container use + container = copy.deepcopy(containers[0]) + container["name"] = "discovery" + + if per_pod_directory: + # Kubeflow multi-replica mode - use pod-specific subdirectories + # Master writes to base dir, workers to worker-{index}/ + # Mount full PVC and create path using downward API labels + container["volumeMounts"] = [ + { + "name": storage_volume_name, + "mountPath": "/data", + } + ] + + # Add replica type and index env vars via downward API + if "env" not in container: + container["env"] = [] + container["env"].extend( + [ + { + "name": "REPLICA_TYPE", + "valueFrom": { + "fieldRef": { + "fieldPath": "metadata.labels['training.kubeflow.org/replica-type']" + } + }, + }, + { + "name": "REPLICA_INDEX", + "valueFrom": { + "fieldRef": { + "fieldPath": "metadata.labels['training.kubeflow.org/replica-index']" + } + }, + }, + ] + ) + + # Modify script to create pod-specific directory based on replica type + # Master: {base_dir}/, Worker: {base_dir}/worker-{index}/ + # SAFETY: base_dir validated earlier in generation, env vars from K8s labels + original_command = ( + container.get("args", [""])[0] if container.get("args") else "" + ) + + # Prepend directory selection logic and update output path + setup_cmd = f""" +if [ "$REPLICA_TYPE" = "master" ]; then + DISCOVERY_DIR=/data/{base_dir} +else + DISCOVERY_DIR=/data/{base_dir}/worker-$REPLICA_INDEX +fi +mkdir -p $DISCOVERY_DIR && cd $DISCOVERY_DIR && """ + + # Replace discovery mount path references with current directory + modified_command = original_command.replace( + f"{DISCOVERY_MOUNT_PATH}/", "./" + ) + # Update final rename command + rename_cmd = f"\nif [ -f ./discovery.json ]; then mv ./discovery.json ./discovery_{sanitized_template_name}.json; fi" + + container["args"] = [setup_cmd + modified_command + rename_cmd] + + else: + # JobSet or single-pod mode - use subPath + # Construct result path + if task_id is not None: + # JobSet mode - task-specific path + result_path = f"{base_dir}/task-{task_id}" + else: + # Single-pod mode - base directory only + result_path = base_dir + + # Update volume mounts to use PVC with subPath + container["volumeMounts"] = [ + { + "name": storage_volume_name, + "mountPath": DISCOVERY_MOUNT_PATH, + "subPath": result_path, + } + ] + + # Add rename command to output discovery_{template_name}.json + # SAFETY: sanitized_template_name validated above to prevent shell injection + original_command = ( + container.get("args", [""])[0] if container.get("args") else "" + ) + rename_cmd = f"\nif [ -f {DISCOVERY_MOUNT_PATH}/discovery.json ]; then mv {DISCOVERY_MOUNT_PATH}/discovery.json {DISCOVERY_MOUNT_PATH}/discovery_{sanitized_template_name}.json; fi" + + if container.get("args"): + container["args"] = [original_command + rename_cmd] + + return container + except (FileNotFoundError, yaml.YAMLError, KeyError, IndexError) as e: + sys.stderr.write( + f"Warning: Could not load discovery template {template_path}: {e}\n" + ) + return None diff --git a/coldpress_common/__init__.py b/coldpress_common/__init__.py index 2c488d0..56678e2 100644 --- a/coldpress_common/__init__.py +++ b/coldpress_common/__init__.py @@ -2,14 +2,16 @@ """Shared validation models and utilities for Coldpress tools.""" from .model import ( - validate_config, + validate_intent, validate_project_config, validate_task_specs, - validate_job_spec, validate_user_config, validate_cluster_config, validate_kubernetes_name, WorkloadConfig, + IntentConfig, + TaskIntent, + DependsOnConfig, ProjectConfig, TaskSpec, JobSpec, @@ -18,14 +20,16 @@ ) __all__ = [ - "validate_config", + "validate_intent", "validate_project_config", "validate_task_specs", - "validate_job_spec", "validate_user_config", "validate_cluster_config", "validate_kubernetes_name", "WorkloadConfig", + "IntentConfig", + "TaskIntent", + "DependsOnConfig", "ProjectConfig", "TaskSpec", "JobSpec", diff --git a/coldpress_common/model.py b/coldpress_common/model.py index bc519de..f23bbb6 100644 --- a/coldpress_common/model.py +++ b/coldpress_common/model.py @@ -39,8 +39,16 @@ early rather than during job execution. """ +import re from typing import Any, Literal, Optional, Union -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import ( + BaseModel, + ConfigDict, + Field, + ValidationError, + field_validator, + model_validator, +) # === Container and Resource Models === @@ -242,8 +250,126 @@ def parse_tasks(cls, v): return v +class DependsOnConfig(BaseModel): + """Task dependency configuration.""" + + task: str + wait_for: Literal["ready", "completion"] + + +class ArgOverride(BaseModel): + """Argument override with optional insertion position.""" + + value: str + insert_after: Optional[str] = None + insert_before: Optional[str] = None + + @model_validator(mode="after") + def validate_insert_position(self): + """Ensure only one of insert_after or insert_before is specified.""" + if self.insert_after and self.insert_before: + raise ValueError( + "Cannot specify both insert_after and insert_before for arg override" + ) + return self + + +class TaskIntent(BaseModel): + """Task intent specification from intent.yaml.""" + + name: str + replicas: Optional[int] = 1 + nodes: Optional[list[int]] = None + args: Optional[dict[str, Union[str, ArgOverride]]] = ( + None # Key-value pairs for arg replacement + ) + env: Optional[dict[str, str]] = None # Key-value pairs for env var replacement + depends_on: Optional[DependsOnConfig] = None + + @field_validator("name") + @classmethod + def validate_task_name(cls, v): + """Validate task name is a valid Kubernetes resource name.""" + return validate_kubernetes_name(v, max_length=63) + + @model_validator(mode="after") + def validate_nodes_match_replicas(self): + """Validate nodes list matches replicas count if specified.""" + if self.nodes is not None and len(self.nodes) != self.replicas: + raise ValueError( + f"Task '{self.name}': nodes list length ({len(self.nodes)}) " + f"must match replicas ({self.replicas})" + ) + return self + + +class IntentConfig(BaseModel): + """Intent.yaml schema - new format for vk8s transformation.""" + + project: str + output: str + target: Literal["jobset", "kubeflow", "kserve", "kuberay"] = ( + "jobset" # Backend to generate + ) + files: Optional[list[str]] = None + discovery: Optional[Union[str, DiscoveryConfig]] = None + tasks: list[TaskIntent] + + @field_validator("discovery", mode="before") + @classmethod + def parse_discovery(cls, v): + """Parse discovery field - accept string or dict.""" + if v is None: + return None + if isinstance(v, str): + return {"template": v, "tasks": "all"} + return v + + @field_validator("tasks") + @classmethod + def validate_tasks_not_empty(cls, v): + """Ensure tasks list is not empty.""" + if not v: + raise ValueError("intent.yaml must include at least one task") + return v + + @model_validator(mode="after") + def validate_task_dependencies(self): + """Validate task dependency references exist.""" + task_names = {task.name for task in self.tasks} + for task in self.tasks: + if task.depends_on and task.depends_on.task not in task_names: + raise ValueError( + f"Task '{task.name}' depends on '{task.depends_on.task}' " + f"which does not exist" + ) + return self + + @model_validator(mode="after") + def validate_no_circular_dependencies(self): + """Validate no circular dependencies exist.""" + # Build dependency graph + deps = { + task.name: task.depends_on.task if task.depends_on else None + for task in self.tasks + } + + # Check each task for cycles + for task_name in deps: + visited = set() + current = task_name + while current is not None: + if current in visited: + raise ValueError( + f"Circular dependency detected involving task '{task_name}'" + ) + visited.add(current) + current = deps.get(current) + return self + + class WorkloadConfig(BaseModel): - """Main config.yaml schema.""" + """Main config.yaml schema (legacy format - for backward compatibility).""" project: Optional[str] = None discovery: Optional[Union[str, DiscoveryConfig]] = None @@ -264,7 +390,7 @@ def parse_discovery(cls, v): @model_validator(mode="after") def validate_has_project(self): - """Ensure project is specified (can be overridden via CLI).""" + """Ensure project is specified (can be overridden via CLI --project flag).""" # Note: project can be None if provided via CLI --project flag return self @@ -281,6 +407,7 @@ class ProjectConfig(BaseModel): """Project configuration from projects/*.yaml.""" namespace: str + targets: Optional[str] = "jobset" # jobset, kubeflow, or comma-separated cluster_queue: Optional[str] = None storage_class: Optional[str] = None storage: Optional[StorageConfig] = None @@ -358,8 +485,6 @@ def validate_kubernetes_name(name: str, max_length: int = 253) -> str: Raises: ValueError: If the name is invalid """ - import re - if not name: raise ValueError("Resource name cannot be empty") @@ -378,22 +503,6 @@ def validate_kubernetes_name(name: str, max_length: int = 253) -> str: return name -def validate_config(config_data: dict) -> WorkloadConfig: - """ - Validate config.yaml data. - - Args: - config_data: Raw YAML data from config file - - Returns: - Validated WorkloadConfig - - Raises: - ValidationError: If validation fails - """ - return WorkloadConfig.model_validate(config_data) - - def validate_project_config(project_data: dict) -> ProjectConfig: """ Validate project configuration. @@ -427,55 +536,55 @@ def validate_task_specs(task_specs_data: list[dict]) -> list[TaskSpec]: for i, task_data in enumerate(task_specs_data): try: validated_tasks.append(TaskSpec.model_validate(task_data)) - except Exception as e: + except (ValidationError, ValueError) as e: # Re-raise with task context raise ValueError(f"Task {i} validation failed: {e}") from e return validated_tasks -def validate_job_spec(job_spec_data: dict) -> JobSpec: +def validate_user_config(user_data: dict) -> UserConfig: """ - Validate complete job specification. + Validate user configuration. Args: - job_spec_data: Complete job spec dict + user_data: Raw YAML data from user file Returns: - Validated JobSpec + Validated UserConfig Raises: ValidationError: If validation fails """ - return JobSpec.model_validate(job_spec_data) + return UserConfig.model_validate(user_data) -def validate_user_config(user_data: dict) -> UserConfig: +def validate_cluster_config(cluster_data: dict) -> ClusterConfig: """ - Validate user configuration. + Validate cluster configuration. Args: - user_data: Raw YAML data from user file + cluster_data: Raw YAML data from cluster file Returns: - Validated UserConfig + Validated ClusterConfig Raises: ValidationError: If validation fails """ - return UserConfig.model_validate(user_data) + return ClusterConfig.model_validate(cluster_data) -def validate_cluster_config(cluster_data: dict) -> ClusterConfig: +def validate_intent(intent_data: dict) -> IntentConfig: """ - Validate cluster configuration. + Validate intent.yaml data. Args: - cluster_data: Raw YAML data from cluster file + intent_data: Raw YAML data from intent file Returns: - Validated ClusterConfig + Validated IntentConfig Raises: ValidationError: If validation fails """ - return ClusterConfig.model_validate(cluster_data) + return IntentConfig.model_validate(intent_data) diff --git a/coldpress_setup/__init__.py b/coldpress_setup/__init__.py index 0c535bd..0933f80 100644 --- a/coldpress_setup/__init__.py +++ b/coldpress_setup/__init__.py @@ -1,4 +1,4 @@ -# Generated by: Claude Sonnet 4.5 +# Assisted by: Claude Sonnet 4.5 """Coldpress Config - Prescriptive cluster configuration.""" -__version__ = "0.2.0" +__version__ = "0.2.1" diff --git a/coldpress_setup/cli.py b/coldpress_setup/cli.py index 696c21e..10ce9e2 100644 --- a/coldpress_setup/cli.py +++ b/coldpress_setup/cli.py @@ -1,4 +1,4 @@ -# Generated by: Claude Sonnet 4.5 +# Assisted by: Claude Sonnet 4.5 """Coldpress Setup CLI - Generate manifests for cluster setup and configuration.""" import click @@ -12,6 +12,7 @@ validate_user_config, validate_cluster_config, ) +from coldpress.constants import COLDPRESS_VERSION from pydantic import ValidationError # Default directories (can be overridden with environment variables) @@ -60,7 +61,7 @@ def _generate_manifest_filename(subcommand, config_file): @click.group() -@click.version_option(version="0.2.0") +@click.version_option(version=COLDPRESS_VERSION) def cli(): """Coldpress Setup - Generate manifests for cluster setup and configuration.""" pass @@ -121,7 +122,14 @@ def cluster(config_file, output_dir): default=None, help=f"Output directory for manifests (default: {MANIFESTS_DIR})", ) -def project(config_file, output_dir): +@click.option( + "--targets", + "-t", + type=str, + default=None, + help="Override target platforms from config (jobset, kubeflow, or comma-separated)", +) +def project(config_file, output_dir, targets): """ Generate project configuration manifests. @@ -151,7 +159,7 @@ def project(config_file, output_dir): click.echo(f"Error: Project config validation failed:\n{e}", err=True) raise SystemExit(1) - raise SystemExit(generate_project_config(config, config_file, output_dir)) + raise SystemExit(generate_project_config(config, config_file, output_dir, targets)) @generate.command() @@ -218,6 +226,9 @@ def generate_cluster_config(config_file, output_dir): for node in validated_config.nodes ] except ValidationError as e: + # Degraded functionality: cluster manifests can still be generated without node-specific + # labeling commands. This allows basic cluster setup even if the cluster config has + # validation errors in the nodes section. Node labeling is an optional enhancement. click.echo(f"Warning: Could not validate cluster config: {e}", err=True) click.echo("Continuing without node labeling commands...", err=True) @@ -234,7 +245,7 @@ def generate_cluster_config(config_file, output_dir): # Copy cluster config to output directory (all documents except the first config one) with open(output_path, "w") as f: # Write only the Kubernetes manifests (skip the config section) - yaml.dump_all(all_docs[1:], f, default_flow_style=False, sort_keys=False) + yaml.safe_dump_all(all_docs[1:], f, default_flow_style=False, sort_keys=False) # Generate node labeling script if nodes were found label_script_path = None @@ -309,7 +320,7 @@ def _generate_and_display_manifests(config, manifest_generator): return manifests -def generate_project_config(config, config_file, output_dir): +def generate_project_config(config, config_file, output_dir, targets_override=None): """Generate project configuration manifests.""" from .generator import generate_project_manifests @@ -318,6 +329,16 @@ def generate_project_config(config, config_file, output_dir): click.echo("Error: Project config must include 'namespace'", err=True) return 1 + # Override targets if specified + if targets_override: + config["targets"] = targets_override + click.echo(f"Using targets from CLI: {targets_override}") + elif "targets" in config: + click.echo(f"Using targets from config: {config['targets']}") + else: + config["targets"] = "jobset" # default + click.echo("Using default targets: jobset") + click.echo("Generating project configuration manifests...") click.echo(f"Namespace: {namespace}") click.echo("\nGenerating resources...") diff --git a/coldpress_setup/generator.py b/coldpress_setup/generator.py index e558147..b7250ee 100644 --- a/coldpress_setup/generator.py +++ b/coldpress_setup/generator.py @@ -1,14 +1,8 @@ -# Generated by: Claude Sonnet 4.5 +# Assisted by: Claude Sonnet 4.5 """Generate Kubernetes manifests for Coldpress cluster setup.""" -import json import yaml - -# Standard labels for all Coldpress-managed resources -COLDPRESS_LABELS = { - "app.kubernetes.io/managed-by": "coldpress", - "app.kubernetes.io/version": "0.2.0", -} +from coldpress.constants import COLDPRESS_LABELS def generate_kueue_resource_flavors(nodes): @@ -20,7 +14,7 @@ def generate_kueue_resource_flavors(nodes): flavors = [] for node_id, node in enumerate(nodes): flavor = { - "apiVersion": "kueue.x-k8s.io/v1beta1", + "apiVersion": "kueue.x-k8s.io/v1beta2", "kind": "ResourceFlavor", "metadata": { "name": f"node{node_id}", @@ -39,22 +33,15 @@ def generate_cluster_queue(name, nodes): Args: name: ClusterQueue name nodes: List of node configs. Node ID is assigned based on position in list. - - Note: - RoCE NIC resources are currently disabled. The roce_nics field in cluster - config is kept for reference but not used in resource generation. """ # Build list of all covered resources across all nodes covered_resources = ["cpu", "memory", "nvidia.com/gpu"] - # NOTE: RoCE NIC resources disabled for now - # Keep roce_nics in cluster config for reference, but don't generate resources - # Build flavors list with resources for each node flavors = [] for node_id, node in enumerate(nodes): - gpus = node.get("gpus", 0) - # roce_nics = node.get("roce_nics", 0) # Kept in config but not used + # Node dicts are created from validated NodeConfig with gpus field guaranteed + gpus = node["gpus"] resources = [ {"name": "cpu", "nominalQuota": "256"}, @@ -62,13 +49,11 @@ def generate_cluster_queue(name, nodes): {"name": "nvidia.com/gpu", "nominalQuota": str(gpus)}, ] - # NOTE: RoCE NIC resources disabled - not added to resources list - flavors.append({"name": f"node{node_id}", "resources": resources}) # Create single resourceGroup with all flavors cluster_queue = { - "apiVersion": "kueue.x-k8s.io/v1beta1", + "apiVersion": "kueue.x-k8s.io/v1beta2", "kind": "ClusterQueue", "metadata": { "name": name, @@ -110,7 +95,7 @@ def generate_namespace(name, storage_size, privileged=False): def generate_local_queue(namespace, cluster_queue_name): """Generate LocalQueue manifest for a namespace.""" local_queue = { - "apiVersion": "kueue.x-k8s.io/v1beta1", + "apiVersion": "kueue.x-k8s.io/v1beta2", "kind": "LocalQueue", "metadata": { "name": f"coldpress-local-queue-{namespace}", @@ -143,8 +128,13 @@ def generate_pvc(name, namespace, storage_size, storage_class): return pvc -def generate_rbac(namespace): - """Generate RBAC manifests for namespace (ServiceAccount, Role, RoleBinding).""" +def generate_rbac(namespace, targets="jobset"): + """Generate RBAC manifests for namespace (ServiceAccount, Role, RoleBinding). + + Args: + namespace: Kubernetes namespace name + targets: Comma-separated list of target platforms (jobset, kubeflow) + """ rbac = [] # ServiceAccount @@ -159,21 +149,56 @@ def generate_rbac(namespace): } rbac.append(sa) - # Role - allow managing JobSets and viewing related resources - role = { - "apiVersion": "rbac.authorization.k8s.io/v1", - "kind": "Role", - "metadata": { - "name": "coldpress-user-role", - "namespace": namespace, - "labels": COLDPRESS_LABELS.copy(), - }, - "rules": [ + # Parse targets + target_list = [t.strip() for t in targets.split(",")] + enable_jobset = "jobset" in target_list + enable_kubeflow = "kubeflow" in target_list + enable_kuberay = "kuberay" in target_list + + # Build RBAC rules based on enabled targets + rules = [] + + # JobSet permissions + if enable_jobset: + rules.append( { "apiGroups": ["jobset.x-k8s.io"], "resources": ["jobsets"], "verbs": ["create", "get", "list", "watch", "delete"], - }, + } + ) + + # Kubeflow Training Operator permissions + if enable_kubeflow: + rules.append( + { + "apiGroups": ["kubeflow.org"], + "resources": ["pytorchjobs", "tfjobs", "mpijobs", "xgboostjobs"], + "verbs": ["create", "get", "list", "watch", "delete"], + } + ) + # KServe permissions + rules.append( + { + "apiGroups": ["serving.kserve.io"], + "resources": ["inferenceservices"], + "verbs": ["create", "get", "list", "watch", "delete"], + } + ) + + # KubeRay permissions + if enable_kuberay: + rules.append( + { + "apiGroups": ["ray.io"], + "resources": ["rayjobs", "rayclusters", "rayservices"], + "verbs": ["create", "get", "list", "watch", "delete"], + } + ) + + # Common permissions (always needed) + rules.extend( + [ { "apiGroups": ["batch"], "resources": ["jobs"], @@ -193,9 +218,29 @@ def generate_rbac(namespace): { "apiGroups": [""], "resources": ["configmaps"], - "verbs": ["create", "get", "list", "watch", "delete"], + "verbs": [ + "create", + "get", + "list", + "watch", + "update", + "patch", + "delete", + ], }, - ], + ] + ) + + # Role - allow managing resources based on enabled targets + role = { + "apiVersion": "rbac.authorization.k8s.io/v1", + "kind": "Role", + "metadata": { + "name": "coldpress-user-role", + "namespace": namespace, + "labels": COLDPRESS_LABELS.copy(), + }, + "rules": rules, } rbac.append(role) @@ -282,42 +327,6 @@ def generate_privileged_scc_binding(namespace): } -def generate_sriov_network_attachments(namespace, roce_nics): - """Generate NetworkAttachmentDefinitions for SRIOV RDMA networks. - - NOTE: This function is currently disabled. RoCE NIC support is kept in - cluster config for reference but NetworkAttachmentDefinitions are not - generated. This function is preserved for future use. - """ - attachments = [] - for nic_id in range(roce_nics): - nic_name = f"eno{5 + nic_id}" - - # Build CNI config using json.dumps to prevent injection - cni_config = { - "cniVersion": "0.3.1", - "name": f"sriov-rdma-net-{nic_name}", - "type": "sriov", - "ipam": {"type": "whereabouts", "range": "192.168.1.0/24"}, - } - - attachment = { - "apiVersion": "k8s.cni.cncf.io/v1", - "kind": "NetworkAttachmentDefinition", - "metadata": { - "name": f"sriov-rdma-net-{nic_name}", - "namespace": namespace, - "labels": COLDPRESS_LABELS.copy(), - "annotations": { - "k8s.v1.cni.cncf.io/resourceName": f"openshift.io/{nic_name}np0rdma" - }, - }, - "spec": {"config": json.dumps(cni_config)}, - } - attachments.append(attachment) - return attachments - - def generate_cluster_manifests(config): """ Generate cluster-wide Kueue manifests (ResourceFlavors, ClusterQueue). @@ -368,7 +377,8 @@ def generate_project_manifests(config): generate_namespace(namespace, storage_size, privileged) ) - # LocalQueue + # LocalQueue (needed for both jobset and kubeflow targets) + # Both JobSet and Kubeflow Training Operator jobs use Kueue for resource management manifests["kueue"].append(generate_local_queue(namespace, cluster_queue_name)) # Storage PVC (for results) @@ -379,7 +389,8 @@ def generate_project_manifests(config): ) # RBAC (ServiceAccount, Role, RoleBinding) - manifests["rbac"].extend(generate_rbac(namespace)) + targets = config.get("targets", "jobset") + manifests["rbac"].extend(generate_rbac(namespace, targets)) return manifests @@ -444,13 +455,6 @@ def generate_all_manifests(config): # Regular user RBAC manifests["rbac"].extend(generate_rbac(ns_name)) - # NOTE: NetworkAttachmentDefinitions for SRIOV disabled - # RoCE NIC support kept in cluster config for reference only - # if max_roce_nics > 0 and not privileged: - # manifests["network"].extend( - # generate_sriov_network_attachments(ns_name, max_roce_nics) - # ) - return manifests @@ -459,8 +463,7 @@ def manifests_to_yaml(manifests): all_docs = [] # Combine in order: kueue, namespaces, storage, rbac - # Note: "network" category removed (RoCE NIC support disabled) for category in ["kueue", "namespaces", "storage", "rbac"]: all_docs.extend(manifests.get(category, [])) - return yaml.dump_all(all_docs, default_flow_style=False, sort_keys=False) + return yaml.safe_dump_all(all_docs, default_flow_style=False, sort_keys=False) diff --git a/discovery/README.md b/discovery/README.md deleted file mode 100644 index 78a6f0f..0000000 --- a/discovery/README.md +++ /dev/null @@ -1,135 +0,0 @@ - -# Discovery Templates - -Pod template for non-privileged user hardware and system discovery. - -## user_snapshot.yaml - -Captures a comprehensive snapshot of the node's system state with performance benchmarks. - -**What it captures:** - -**System Information:** -- CPU model, count, cores per socket, threads per core -- CPU frequency (current, max, min) -- CPU cache (L1d, L1i, L2, L3) -- Memory (total, available, free) -- NUMA node count -- Kernel version and uptime - -**GPU Information:** -- GPU count and models (node-level info) -- GPU memory per device -- GPU driver and CUDA version - -**Network Information:** -- Network interface names -- RDMA devices and count -- MTU settings -- TCP buffer sizes and congestion control - -**Storage Information:** -- Block devices and NVMe count -- tmpfs size and availability -- I/O scheduler configuration - -**Performance Configuration:** -- CPU governor -- NUMA balancing -- Transparent hugepages -- Swappiness and dirty ratios -- File and process limits - -**Process Limits:** -- Max user processes and open files -- Stack and memory size limits -- Cgroup CPU quota and memory limits - -**Performance Benchmarks:** -- CPU compute (MFLOPS) -- Memory bandwidth (read/write) -- Disk I/O (read/write) -- Context switch overhead - -**Requirements:** -- Non-privileged (no special permissions) -- Uses PyTorch image for Python benchmarks - -**Output:** -Single JSON file: `/tmp/result/discovery.json` - -## Usage - -Discovery templates are Pod specs that get stitched into JobSets by the coldpress generator. - -### In Job Spec - -```yaml -name: my-job -namespace: coldpress-project - -discovery: - template: discovery_templates/user_snapshot.yaml - -tasks: - - name: my-workload - image: my-image:latest - args: ["--arg1", "value"] - gpus: 1 -``` - -The generator will: -1. Read the discovery template Pod spec -2. Add it as a discovery task in the JobSet (runs first) -3. Set completion blocking (main tasks wait for discovery) -4. Extract results to PVC at `coldpress_results//discovery/snapshot.txt` - -## Output Example - -```json -{ - "timestamp": "2026-04-08T16:30:00.123456", - "system": { - "cpu_model": "Intel(R) Xeon(R) Platinum 8358 CPU @ 2.60GHz", - "cpu_count": "128", - "cpu_cores_per_socket": "64", - "cpu_threads_per_core": "2", - "memory_total_gb": "512.00", - "cache_l3": "48M", - "numa_node_count": "2" - }, - "gpu": { - "gpu_count": "8", - "gpu_models": "NVIDIA H100 PCIe", - "gpu_memory_per_device_mb": "81559", - "gpu_driver_version": "535.129.03", - "cuda_version": "12.2" - }, - "network": { - "interface_names": "eth0, eno5np0", - "rdma_devices": "mlx5_0, mlx5_1", - "rdma_device_count": "2" - }, - "benchmarks": { - "cpu_compute": { - "iterations": 10000000, - "time_seconds": "0.8234", - "mflops": "24.31" - }, - "memory_bandwidth": { - "write_bandwidth_mbps": "3456.78", - "read_bandwidth_mbps": "4123.45" - }, - "disk_io": { - "write_bandwidth_mbps": "523.12", - "read_bandwidth_mbps": "678.34" - } - } -} -``` - -## Non-Privileged Only - -This template runs without privileged access and captures only user-visible state. - -For privileged discovery (PCIe topology, RDMA devices, hardware introspection), see `legacy/examples/discovery/`. diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md deleted file mode 100644 index 8bf2b9e..0000000 --- a/docs/CHANGELOG.md +++ /dev/null @@ -1,151 +0,0 @@ -# Changelog - -All notable changes to Coldpress will be documented in this file. - -## [Unreleased] - -### Added -- **Comprehensive test suite with CI/CD** - Automated testing on every commit - - Organized all tests in `tests/` directory - - Created `tests/run_all_tests.sh` to run full test suite locally - - GitHub Actions workflow runs tests on Python 3.9, 3.10, 3.11, 3.12, 3.13, 3.14 - - Tests cover: validation, labels, security, error handling, RoCE disabled, namespace consistency - - All tests pass on multiple Python versions - - Validates fixes for GitHub issue #37 - -- **Security and input validation improvements** - Prevent injection attacks and validate user input - - JSON construction now uses `json.dumps()` instead of f-strings in `generate_sriov_network_attachments()` - - Added `validate_kubernetes_name()` function to validate resource names against Kubernetes spec - - Pydantic validators now enforce Kubernetes naming rules (lowercase, alphanumeric, dashes, dots only) - - Task names, namespace names, and usernames validated at load time - - Prevents injection attacks via malformed resource names - - No hardcoded temp file paths (timestamped manifests used instead) - - Resolves GitHub issue #37 (Security and Reliability Issues) - -- **Improved error handling** - More specific exception handling and proper exit code propagation - - Replaced overly broad `except Exception` with specific exception types (FileNotFoundError, yaml.YAMLError, KeyError, etc.) - - CLI commands now properly propagate exit codes to shell using `raise SystemExit(code)` - - Exit code 0: Success, 1: Application errors, 2: Usage/validation errors - - Discovery template loading errors provide specific error messages - - Service creation errors are logged with details - - Resolves GitHub issue #37 (Error Handling) - -- **Standard resource labels** - All generated resources now include standard Kubernetes labels for easy querying and management - - `app.kubernetes.io/managed-by: coldpress` - Identifies all Coldpress-managed resources - - `app.kubernetes.io/version: 0.2.0` - Tracks Coldpress version - - `coldpress.io/job-id: {job_name}` - Job-specific identifier for compute resources - - Enables single-command queries: `kubectl get all -A -l app.kubernetes.io/managed-by=coldpress` - - Enables bulk operations: `kubectl delete all -n namespace -l coldpress.io/job-id=job-name` - - See `docs/LABELS.md` for detailed usage examples - - Resolves GitHub issue #37 (No Consistent Resource Labelling) - -- **Timestamped node labeling scripts** - Node labeling scripts now include timestamps in filename (e.g., `label-nodes-cluster-20260413-152928.sh`) for consistency with manifests - -- **Per-task hardware discovery** - Discovery now runs as init container for each task instead of a separate job - - Captures hardware info on the actual node where each task executes (not a different node) - - Discovery results placed in task-specific directories: `{base_dir}/task-{task_id}/discovery_{template}.json` - - New discovery config format supports selecting which tasks run discovery: - - Simple format: `discovery: user_snapshot` (all tasks, backward compatible) - - Detailed format: `discovery: {template: user_snapshot, tasks: all}` or `tasks: [0, 1]` - - Solves the dynamic scheduling problem: discovery reflects actual hardware used by each task - - Mkdir job now creates task subdirectories: `task-0/`, `task-1/`, etc. - -### Fixed -- **Namespace generation duplication removed** - `generate_project_manifests()` now uses shared `generate_namespace()` function - - Eliminates duplicated namespace creation logic in `coldpress_setup/generator.py` - - Ensures consistent labeling across all code paths (both use `COLDPRESS_LABELS`) - - Privileged namespaces now properly supported in project manifests (security labels applied correctly) - - Added `tests/test_namespace_consistency.py` to verify no duplication exists - - Resolves GitHub issue #37 (Duplicated Logic) - -### Changed -- **`coldpress-setup` now generates manifests instead of applying directly** - - Commands now write timestamped manifest files to `manifests/` directory (configurable with `--output-dir`) - - Admin is responsible for reviewing and applying manifests with `oc apply -f` - - Manifest filenames include subcommand, config name, and timestamp (e.g., `project-coldpress-project-20260413-152928.yaml`) - - All RBAC permissions and resources are now visible in generated manifests before application - - Enables GitOps workflows and better audit trails - - Removed `--dry-run` flag (no longer needed) - - Removed `--output` flag (replaced with `--output-dir`) - - Removed `COLDPRESS_OC_FLAGS` environment variable (CLI no longer interacts with cluster) - - Added `COLDPRESS_MANIFESTS_DIR` environment variable (default: `manifests/`) - -- **`coldpress` manifest generation improvements** - - Removed automatic cluster-based node allocation - - Default behavior: Kubernetes scheduler selects any node with `coldpress.node` label - - Node assignment priority: CLI `--node` flag > `nodes` in config.yaml > Kubernetes scheduler - - Can specify nodes in config.yaml: `nodes: [0, 1]` or via CLI: `--node 0 --node 1` - - Generated manifests use `nodeAffinity` (scheduler picks node) or `nodeSelector` (pinned to specific node) - -### Removed -- **RoCE NIC support temporarily disabled** - Simplified cluster configuration - - No longer generates RDMA resources in ClusterQueue (`openshift.io/eno*np0rdma`) - - No longer generates NetworkAttachmentDefinitions for SRIOV - - The `roce_nics` field is kept in cluster config for reference/future use - - Simplifies cluster setup while keeping the option to re-enable later - - Functions preserved but disabled: `generate_sriov_network_attachments()` - -- **`coldpress-setup verify` command** - Admins can verify resources using standard kubectl/oc commands - - Use `oc get namespaces`, `oc get pvc -n `, etc. instead - -- **`coldpress/allocator.py` module** - Removed automatic node allocation - - Previous behavior queried cluster to find least-loaded node - - New behavior: Kubernetes scheduler handles allocation - - Users can pin tasks to specific nodes with `--node` flag - -## [0.2.0] - 2026-04-13 - -### Changed -- **Version downgrade from 2.0.0 to 0.2.0** - Reset version to reflect early development stage -- **Switched to `uv` package manager** - - `setup-env.sh` now uses `uv` for virtual environment and package installation - - Virtual environment renamed from `venv/` to `.venv/` -- **Removed `kubernetes` Python dependency** - Replaced with direct `oc`/`kubectl` subprocess calls - - Better compatibility across different cluster configurations - - Applies to `coldpress/allocator.py` and `coldpress_setup/cli.py` -- **Refactored `coldpress-setup` CLI** - Changed from single command to subcommand structure: - - `coldpress-setup apply ` → `coldpress-setup generate cluster|project|user ` - - Auto-resolves config files from standard directories (cluster/, projects/, users/) - - Commands work with just filename: `coldpress-setup generate cluster ocp-test.yaml` - -### Added -- **Shared validation module** (`coldpress_common/model.py`) - Config files are now validated at load time by both tools - - Validates config.yaml, job-spec.yaml, project configs, and user configs - - Clear error messages when YAML structure is invalid - - Both `coldpress` and `coldpress-setup` validate before operations - - New `UserConfig` model for validating users/*.yaml files -- **`pyproject.toml`** - Modern Python packaging configuration -- **`docs/VALIDATION.md`** - Documentation for config validation -- **`test_validation.py`** - Test script for validating Pydantic models -- **kubectl/oc requirement check** - `setup-env.sh` now verifies kubectl or oc is installed - -### Improved -- **`coldpress` CLI refactoring** - Extracted helper functions for better code organization: - - `_load_and_validate_config()` - Config loading with validation - - `_load_task_specs()` - Task spec loading with validation - - `_allocate_nodes_for_tasks()` - Node allocation logic - - `_prepare_configmap_files()` - ConfigMap preparation - - `_write_output_files()` - Output file generation -- **`coldpress-setup` CLI refactoring** - Modular functions for each operation: - - `_get_kubectl_cmd()` - Auto-detect kubectl or oc - - `_get_kubectl_flags()` - Load flags from environment variable - - `_apply_yaml_to_cluster()` - Apply YAML via stdin - - `_verify_project()` - Verify project resources exist -- **`coldpress/generator.py` refactoring** - Extracted helper functions: - - `_infer_blocking_type_and_health_check()` - Infer task blocking behavior - - `_substitute_dns_in_args()` - DNS placeholder substitution - - `_build_init_jobs()` - Build initialization jobs -- **Documentation updates**: - - `steps.md` - Comprehensive end-to-end validation guide - - `README.md` - Updated commands to reflect new CLI structure - -### Removed -- **`TODO.md`** - Removed temporary task tracking file - -### Fixed -- Permission issues when running without admin privileges - Now properly handled via environment variable -- Config file path resolution - Automatically checks standard directories - -## [Previous Versions] - -See git history for changes prior to 0.2.0 diff --git a/docs/ERROR_HANDLING.md b/docs/ERROR_HANDLING.md deleted file mode 100644 index ce5b576..0000000 --- a/docs/ERROR_HANDLING.md +++ /dev/null @@ -1,202 +0,0 @@ -# Error Handling - -Coldpress implements robust error handling with clear error messages and proper exit codes. - -## Exit Codes - -All Coldpress CLI commands return standard Unix exit codes: - -| Exit Code | Meaning | Example | -|-----------|---------|---------| -| 0 | Success | Job generated successfully | -| 1 | Application error | File not found, validation failed, template error | -| 2 | Usage error | Missing required option, invalid argument | - -### Testing Exit Codes - -```bash -# Test success (exit 0) -coldpress generate --config examples/pytorch_ddp_training/config.yaml -echo $? # Should print: 0 - -# Test error (exit 1) -coldpress-setup generate cluster non-existent.yaml -echo $? # Should print: 1 - -# Test usage error (exit 2) -coldpress generate -echo $? # Should print: 2 -``` - -## Exception Handling - -### Specific Exception Types - -Coldpress uses specific exception types instead of broad `except Exception` blocks: - -```python -# ✅ Good: Specific exceptions -try: - with open(template_path, "r") as f: - template = yaml.safe_load(f) -except (FileNotFoundError, yaml.YAMLError, KeyError, IndexError) as e: - print(f"Warning: Could not load template {template_path}: {e}") - return None - -# ❌ Bad: Overly broad exception -try: - with open(template_path, "r") as f: - template = yaml.safe_load(f) -except Exception as e: # Catches too much - return None -``` - -### Exception Categories - -**File Operations:** -- `FileNotFoundError` - File or directory doesn't exist -- `PermissionError` - Insufficient permissions - -**YAML Parsing:** -- `yaml.YAMLError` - Invalid YAML syntax -- `yaml.scanner.ScannerError` - YAML scanning errors - -**Validation:** -- `pydantic.ValidationError` - Schema validation failed -- `ValueError` - Invalid value or configuration -- `KeyError` - Required key missing from dict -- `IndexError` - List index out of range - -**URL/Service:** -- `ValueError` - URL parsing failed -- `AttributeError` - Missing attribute in object -- `KeyError` - Service configuration incomplete - -## Error Messages - -Coldpress provides clear, actionable error messages: - -### Configuration Errors - -```bash -$ coldpress-setup generate project invalid-project.yaml -Error: Project config validation failed: -1 validation error for ProjectConfig -namespace - Field required [type=missing, input_value={...}, input_type=dict] -``` - -### File Not Found - -```bash -$ coldpress generate --config missing.yaml -Error: Config file not found: missing.yaml -``` - -### Discovery Template Errors - -```bash -Warning: Could not load discovery template discovery/invalid.yaml: - while scanning for the next token - found character '\t' that cannot start any token -``` - -## Validation - -Coldpress validates configurations at load time using Pydantic models: - -```python -from coldpress_common import validate_config, validate_project_config -from pydantic import ValidationError - -try: - config = validate_config(config_data) -except ValidationError as e: - print(f"Validation failed: {e}") - sys.exit(1) -``` - -This catches errors **before** attempting to generate manifests, saving time and providing clear feedback. - -## Best Practices - -### 1. Check Exit Codes in Scripts - -```bash -#!/bin/bash -set -e # Exit on any error - -coldpress-setup generate cluster cluster/prod.yaml -if [ $? -ne 0 ]; then - echo "Cluster setup failed!" - exit 1 -fi - -coldpress generate --config job.yaml -if [ $? -ne 0 ]; then - echo "Job generation failed!" - exit 1 -fi -``` - -### 2. Validate Before Deploying - -```bash -# Generate manifests (validates config) -coldpress-setup generate project projects/my-project.yaml - -# Review generated manifests -cat manifests/project-my-project-*.yaml - -# Apply only after manual review -oc apply -f manifests/project-my-project-*.yaml -``` - -### 3. Handle Errors Gracefully - -```python -import subprocess -import sys - -result = subprocess.run( - ["coldpress", "generate", "--config", "job.yaml"], - capture_output=True, - text=True -) - -if result.returncode != 0: - print(f"Error: {result.stderr}") - sys.exit(result.returncode) -``` - -## Debugging - -### Enable Traceback - -For unexpected errors, Coldpress prints a full traceback: - -```bash -$ coldpress generate --config bad.yaml -Error generating JobSet: division by zero -Traceback (most recent call last): - File "coldpress/cli.py", line 365, in generate - ... -ZeroDivisionError: division by zero -``` - -### Verbose YAML Errors - -YAML parsing errors show line and column numbers: - -```bash -Warning: Could not load discovery template discovery/bad.yaml: - while parsing a block mapping - in "discovery/bad.yaml", line 5, column 3 - expected , but found '' - in "discovery/bad.yaml", line 6, column 4 -``` - -## Related - -- [Validation](VALIDATION.md) - Pydantic model validation -- [Labels](LABELS.md) - Resource labels for querying and cleanup diff --git a/docs/LABELS.md b/docs/LABELS.md deleted file mode 100644 index b13776d..0000000 --- a/docs/LABELS.md +++ /dev/null @@ -1,100 +0,0 @@ -# Coldpress Resource Labels - -All resources generated by Coldpress tools are labeled with standard Kubernetes labels for easy identification and management. - -## Standard Labels - -Every resource created by `coldpress` or `coldpress-setup` includes these labels: - -```yaml -labels: - app.kubernetes.io/managed-by: coldpress - app.kubernetes.io/version: "0.2.0" -``` - -### Job-specific Labels - -Resources created by `coldpress generate` also include: - -```yaml -labels: - coldpress.io/job-id: "{job_name}" -``` - -## Querying Resources - -### Find all Coldpress-managed resources - -```bash -# All resources in a specific namespace -kubectl get all,pvc,configmap,rolebinding -n my-namespace -l app.kubernetes.io/managed-by=coldpress - -# All resources across all namespaces -kubectl get all,pvc,configmap,rolebinding -A -l app.kubernetes.io/managed-by=coldpress - -# Cluster-wide resources (ResourceFlavors, ClusterQueues) -kubectl get resourceflavors,clusterqueues -l app.kubernetes.io/managed-by=coldpress -``` - -### Find resources for a specific job - -```bash -# Find all resources for a specific job -kubectl get all,configmap -n my-namespace -l coldpress.io/job-id=pytorch-training - -# Get job details -kubectl describe jobset -n my-namespace -l coldpress.io/job-id=pytorch-training -``` - -### Find resources by version - -```bash -# Find all resources created by Coldpress v0.2.0 -kubectl get all -A -l app.kubernetes.io/version=0.2.0 -``` - -## Cleanup Operations - -### Delete all Coldpress resources from a namespace - -```bash -# Delete all compute resources (JobSets, Services, ConfigMaps) -kubectl delete jobset,service,configmap -n my-namespace -l app.kubernetes.io/managed-by=coldpress - -# Delete a specific job -kubectl delete jobset,service,configmap -n my-namespace -l coldpress.io/job-id=pytorch-training -``` - -### Audit Coldpress resources - -```bash -# List all namespaces managed by Coldpress -kubectl get namespaces -l app.kubernetes.io/managed-by=coldpress - -# Count Coldpress resources -kubectl get all -A -l app.kubernetes.io/managed-by=coldpress --no-headers | wc -l - -# Export all Coldpress resources for backup -kubectl get all,pvc,configmap,rolebinding -A -l app.kubernetes.io/managed-by=coldpress -o yaml > coldpress-backup.yaml -``` - -## Resource Types with Labels - -### coldpress-setup resources: -- **Cluster-wide:** ResourceFlavor, ClusterQueue -- **Namespace-scoped:** Namespace, LocalQueue, PersistentVolumeClaim, ServiceAccount, Role, RoleBinding, NetworkAttachmentDefinition - -### coldpress resources: -- **Job-scoped:** JobSet, Service, ConfigMap - -## Label Schema - -Coldpress follows the [Kubernetes Recommended Labels](https://kubernetes.io/docs/concepts/overview/working-with-objects/common-labels/) standard: - -| Label | Description | Example | -|-------|-------------|---------| -| `app.kubernetes.io/managed-by` | Tool managing the resource | `coldpress` | -| `app.kubernetes.io/version` | Version of Coldpress | `0.2.0` | -| `coldpress.io/job-id` | Job identifier (jobs only) | `pytorch-training` | - -Additional labels may be added in future versions for finer-grained control. diff --git a/docs/ROCE_DISABLED.md b/docs/ROCE_DISABLED.md deleted file mode 100644 index eee34b3..0000000 --- a/docs/ROCE_DISABLED.md +++ /dev/null @@ -1,127 +0,0 @@ -# RoCE NIC Support - Temporarily Disabled - -## Overview - -RoCE (RDMA over Converged Ethernet) NIC support has been temporarily disabled in Coldpress to simplify cluster configuration. The `roce_nics` field is preserved in cluster config files for reference and potential future re-enablement. - -## What Changed - -### Disabled Features - -1. **ClusterQueue RDMA Resources** - No longer generates resources like: - - `openshift.io/eno5np0rdma` - - `openshift.io/eno6np0rdma` - - etc. - -2. **NetworkAttachmentDefinitions** - No longer generates SRIOV network attachments for RDMA - -3. **Network Category in Manifests** - The `network` manifest category is no longer used - -### What's Preserved - -1. **`roce_nics` Field in Config** - The field still exists in: - - `coldpress_common.model.NodeConfig` - - `coldpress_common.model.TaskSpec` - - Cluster YAML config files - -2. **Generator Functions** - Functions are preserved but disabled: - - `generate_sriov_network_attachments()` - Marked as disabled in docstring - -## Configuration Files - -### Cluster Config (cluster/*.yaml) - -You can still include `roce_nics` in node specifications: - -```yaml -nodes: - - hostname: node1 - gpus: 2 - roce_nics: 2 # ← Still accepted, but ignored during manifest generation - - hostname: node2 - gpus: 4 - roce_nics: 1 # ← Still accepted, but ignored during manifest generation -``` - -The field is validated and stored but **not used** when generating manifests. - -### Generated Resources - -**Before (with RoCE enabled):** -```yaml -apiVersion: kueue.x-k8s.io/v1beta1 -kind: ClusterQueue -spec: - resourceGroups: - - coveredResources: - - cpu - - memory - - nvidia.com/gpu - - openshift.io/eno5np0rdma # ← RoCE resources - - openshift.io/eno6np0rdma # ← RoCE resources -``` - -**After (RoCE disabled):** -```yaml -apiVersion: kueue.x-k8s.io/v1beta1 -kind: ClusterQueue -spec: - resourceGroups: - - coveredResources: - - cpu - - memory - - nvidia.com/gpu # ← Only GPU resources -``` - -## Benefits of Disabling - -1. **Simpler Setup** - Cluster admins don't need to configure SRIOV or RDMA -2. **Fewer Dependencies** - No need for SR-IOV CNI or NetworkAttachmentDefinition CRDs -3. **Cleaner Manifests** - Fewer resources to review and apply -4. **Easier Debugging** - Less complexity when troubleshooting networking issues - -## Re-enabling RoCE Support - -If you need RoCE support in the future, you can re-enable it by: - -1. Uncommenting the code in `coldpress_setup/generator.py`: - - Line ~412: `max_roce_nics = max([node.get("roce_nics", 0) for node in nodes], default=0)` - - Line ~419: `"network": []` in manifests dict - - Line ~458-461: NetworkAttachmentDefinitions generation - - Lines in `generate_cluster_queue()`: RoCE resource addition - -2. Update `manifests_to_yaml()` to include `"network"` in the category list - -3. Remove the "disabled" notes from function docstrings - -4. Run tests to verify: `python test_roce_disabled.py` - -## Code Locations - -**Disabled Code:** -- `coldpress_setup/generator.py:generate_cluster_queue()` - Lines ~35-65 -- `coldpress_setup/generator.py:generate_all_manifests()` - Lines ~411-461 -- `coldpress_setup/generator.py:generate_sriov_network_attachments()` - Entire function - -**Preserved Config:** -- `coldpress_common/model.py:NodeConfig.roce_nics` - Line ~203 -- `coldpress_common/model.py:TaskSpec.roce_nics` - Line ~158 - -## Testing - -Run the RoCE disabled test suite: - -```bash -python test_roce_disabled.py -``` - -This verifies: -- ✅ No RoCE resources in ClusterQueue -- ✅ No NetworkAttachmentDefinitions generated -- ✅ `roce_nics` field still validates in config -- ✅ Existing tests still pass - -## Related - -- [Security](../README.md#security) - Input validation improvements -- [Labels](LABELS.md) - Resource labeling for queries diff --git a/docs/TESTING.md b/docs/TESTING.md deleted file mode 100644 index cd67ad6..0000000 --- a/docs/TESTING.md +++ /dev/null @@ -1,228 +0,0 @@ -# Coldpress Testing Guide - -This document describes the comprehensive test suite for Coldpress and how to use it. - -## Test Organization - -All tests are located in the `tests/` directory: - -``` -tests/ -├── README.md # Detailed test documentation -├── run_all_tests.sh # Run all tests locally -├── test_validation.py # Pydantic model validation -├── test_labels.py # Resource labels (issue #37-8) -├── test_security.py # Security & input validation (issue #37-6) -├── test_error_handling.py # Error handling (issue #37-7) -├── test_exit_codes.sh # Shell exit codes -└── test_roce_disabled.py # RoCE NIC disabled tests -``` - -## Running Tests Locally - -### Run All Tests - -```bash -./tests/run_all_tests.sh -``` - -This runs all 6 test suites and provides a summary: - -``` -════════════════════════════════════════════════════════════ - TEST SUMMARY -════════════════════════════════════════════════════════════ - - ✅ Passed: 6 - ❌ Failed: 0 - - 🎉 All tests passed! -════════════════════════════════════════════════════════════ -``` - -### Run Individual Tests - -```bash -# Validation tests -python tests/test_validation.py - -# Label tests -python tests/test_labels.py - -# Security tests -python tests/test_security.py - -# Error handling tests -python tests/test_error_handling.py - -# RoCE disabled tests -python tests/test_roce_disabled.py - -# Exit code tests (shell script) -bash tests/test_exit_codes.sh -``` - -## Continuous Integration - -### GitHub Actions - -Tests run automatically on: -- Every push to `main` or `v0.2` branches -- Every pull request to `main` or `v0.2` branches - -**Workflow:** `.github/workflows/tests.yml` - -**Python versions tested:** -- Python 3.9 -- Python 3.10 -- Python 3.11 -- Python 3.12 -- Python 3.13 -- Python 3.14 - -**View results:** -- Check the Actions tab on GitHub -- Badge on README shows latest status: [![Tests](https://github.com/asanaullah/coldpress/workflows/Tests/badge.svg)](https://github.com/asanaullah/coldpress/actions) - -### What Gets Tested - -1. **Validation Tests** - Verifies Pydantic models catch errors - - Config validation (config.yaml) - - Task spec validation (job-spec.yaml) - - Project config validation - - User config validation - - Fixes: Issue #37-1 (No Input Validation) - -2. **Label Tests** - Verifies resource labeling - - All resources have `app.kubernetes.io/managed-by: coldpress` - - Version labels present - - Job-specific labels on job resources - - Fixes: Issue #37-8 (No Consistent Resource Labelling) - -3. **Security Tests** - Verifies security improvements - - JSON constructed safely (no f-string injection) - - Kubernetes names validated - - User input sanitized - - No hardcoded temp files - - Fixes: Issue #37-6 (Security and Reliability Issues) - -4. **Error Handling Tests** - Verifies error handling - - Exit codes propagate correctly (0/1/2) - - Specific exception types used - - Clear error messages - - Validation errors caught early - - Fixes: Issue #37-7 (Error Handling) - -5. **Exit Code Tests** - Verifies shell-level exit codes - - Success returns 0 - - Errors return 1 - - Usage errors return 2 - -6. **RoCE Disabled Tests** - Verifies RoCE NIC support disabled - - No RDMA resources in ClusterQueue - - No NetworkAttachmentDefinitions generated - - `roce_nics` field still validates - -## Test Dependencies - -**Required packages:** -```bash -pip install pyyaml click pydantic -``` - -**Or use the development environment:** -```bash -source .venv/bin/activate # or setup-env.sh -``` - -## Adding New Tests - -When adding new functionality: - -1. **Create test file:** `tests/test_.py` -2. **Add to test runner:** Edit `tests/run_all_tests.sh` -3. **Add to CI:** Edit `.github/workflows/tests.yml` -4. **Document:** Update `tests/README.md` - -### Test Template - -```python -#!/usr/bin/env python3 -"""Test description.""" - -import sys - -def test_feature(): - """Test specific feature.""" - print("=" * 60) - print("Testing Feature") - print("=" * 60) - - # Test logic here - if condition: - print("✅ Test passed") - return True - else: - print("❌ Test failed") - return False - -def main(): - """Run all tests.""" - if test_feature(): - print("\n✅ All tests passed!") - return 0 - else: - print("\n❌ Some tests failed") - return 1 - -if __name__ == "__main__": - sys.exit(main()) -``` - -## Debugging Test Failures - -### Local Debugging - -1. **Run specific test:** - ```bash - python tests/test_.py - ``` - -2. **Check error output:** - Tests print detailed error messages showing what failed - -3. **Run with verbose output:** - Most tests print progress as they run - -### CI Debugging - -1. **View Actions logs:** - - Go to GitHub repository - - Click "Actions" tab - - Click on failed workflow run - - Expand failed step - -2. **Common issues:** - - Import errors → Check `PYTHONPATH` in workflow - - Missing dependencies → Check `pip install` step - - Python version issues → Check matrix strategy - -## Test Coverage - -Current test coverage addresses all major categories from GitHub issue #37: - -| Category | Issue # | Test File | Status | -|----------|---------|-----------|--------| -| Input Validation | #37-1 | test_validation.py | ✅ | -| Separation of Concerns | #37-2 | Integration tests | ✅ | -| Error Handling | #37-3, #37-7 | test_error_handling.py | ✅ | -| Security | #37-6 | test_security.py | ✅ | -| Resource Labeling | #37-8 | test_labels.py | ✅ | -| RoCE Disabled | N/A | test_roce_disabled.py | ✅ | - -## Related Documentation - -- [tests/README.md](tests/README.md) - Detailed test documentation -- [docs/ERROR_HANDLING.md](docs/ERROR_HANDLING.md) - Error handling guide -- [docs/LABELS.md](docs/LABELS.md) - Resource labels guide -- [docs/ROCE_DISABLED.md](docs/ROCE_DISABLED.md) - RoCE NIC disabled guide diff --git a/docs/VALIDATION.md b/docs/VALIDATION.md deleted file mode 100644 index a365f43..0000000 --- a/docs/VALIDATION.md +++ /dev/null @@ -1,215 +0,0 @@ -# Coldpress YAML Validation - -Coldpress includes comprehensive YAML validation using Pydantic models. This ensures that configuration errors are caught early at file load time, rather than during job execution. - -## Overview - -The validation system is implemented in the shared `coldpress_common` module and provides type-safe, schema-validated models for all Coldpress configuration files used by both `coldpress` and `coldpress-setup` tools: - -- **config.yaml** - Workload configuration (`coldpress` tool) -- **job-spec.yaml** - Task specifications (`coldpress` tool) -- **projects/*.yaml** - Project configuration (both tools) -- **users/*.yaml** - User configuration (`coldpress-setup` tool) - -## How It Works - -Both `coldpress` and `coldpress-setup` tools validate YAML files automatically: - -### `coldpress generate` -1. **Config validation** - Checks project, discovery, output, and files fields -2. **Task spec validation** - Validates container specs, resources, volumes, and dependencies -3. **Project config validation** - Ensures namespace and storage are properly configured - -### `coldpress-setup generate` -1. **Project config validation** - Validates namespace, storage, cluster_queue -2. **User config validation** - Validates username and namespaces list - -If validation fails, you'll get clear error messages indicating exactly what's wrong and the manifest generation will be aborted. - -## Example Error Messages - -### Missing required field: -``` -Error: Task spec validation failed: 1 validation error for TaskSpec -name - Field required [type=missing, input_value={'containers': [...]}, input_type=dict] -``` - -### Invalid blocking configuration: -``` -Error: Task 'my-task' has blocking='endpoint' but no health_check or readinessProbe -``` - -### Invalid resource specification: -``` -Error: Task spec validation failed: 1 validation error for TaskSpec -containers.0.resources.requests - Field required [type=missing] -``` - -## Testing Validation - -Run the validation test script to verify the system is working: - -```bash -./test_validation.py -``` - -This will test both valid and invalid configurations to ensure proper error detection. - -## Using Validation in Code - -You can import and use the validation functions directly from the shared module: - -```python -import yaml -from coldpress_common import ( - validate_config, - validate_task_specs, - validate_project_config, - validate_user_config, -) -from pydantic import ValidationError - -# Validate config.yaml -try: - with open("config.yaml") as f: - config_data = yaml.safe_load(f) - config = validate_config(config_data) - print(f"Valid config: {config.project}") -except ValidationError as e: - print(f"Validation error: {e}") - -# Validate user config (file is in users/ directory) -try: - with open("users/username.yaml") as f: - user_data = yaml.safe_load(f) - user = validate_user_config(user_data) - print(f"Valid user: {user.username}") -except ValidationError as e: - print(f"Validation error: {e}") -``` - -## Supported Validations - -### Config File (`config.yaml`) -- ✓ project (optional, can be overridden via CLI) -- ✓ discovery (optional) -- ✓ output (optional) -- ✓ files (optional list of strings) -- ✓ nodes (optional list of integers) - explicit node assignments for tasks - -### Task Spec (`job-spec.yaml`) -**Required:** -- ✓ name (string) -- ✓ containers (list with at least one container) - -**Optional:** -- ✓ blocking ("completion" or "endpoint") -- ✓ health_check (required if blocking="endpoint") -- ✓ resources (requests and limits) -- ✓ volumes (emptyDir, PVC, configMap) -- ✓ env (environment variables) -- ✓ tolerate_all (boolean) -- ✓ network_mode ("host" or "default") -- ✓ privileged (boolean) -- ✓ sys_mounts (host path mounts) - -**Container-level:** -- ✓ name, image (required) -- ✓ command, args (optional) -- ✓ workingDir (optional) -- ✓ resources (optional) -- ✓ env (optional) -- ✓ ports (optional) -- ✓ readinessProbe (optional) - -### Project Config (stored in projects/ directory) -**Required:** -- ✓ namespace (string) - -**Optional:** -- ✓ cluster_queue (string) -- ✓ storage_class (string) -- ✓ storage (object with results, models, size) - -### User Config (stored in users/ directory) -**Required:** -- ✓ username (string) -- ✓ namespaces (list of strings, at least one required) - -**Validation Rules:** -- namespaces list cannot be empty - -## Advanced Validation Rules - -1. **Endpoint Blocking**: If `blocking: "endpoint"`, task must have either: - - A `health_check` URL, OR - - A `readinessProbe` in the first container - -2. **Resource Consistency**: GPU resources are automatically tracked and validated - -3. **Volume References**: Volume names are validated for internal consistency - -4. **Environment Variables**: Can be specified as dict or list format - -## Benefits - -- **Early Error Detection**: Catch configuration errors before job submission -- **Clear Error Messages**: Pydantic provides detailed, actionable error messages -- **Type Safety**: Ensures fields have correct types (string, int, list, etc.) -- **Schema Documentation**: Pydantic models serve as living documentation -- **IDE Support**: Better autocomplete and type hints when working with configs - -## Architecture - -### Shared Validation Module - -The validation logic is centralized in `coldpress_common/model.py`, which is imported by both: -- `coldpress/cli.py` - Workload generation tool -- `coldpress_setup/cli.py` - Cluster setup tool - -This ensures: -- **Consistency**: Same validation rules across both tools -- **Maintainability**: Single source of truth for schemas -- **Early detection**: Errors caught before cluster operations - -### Why Shared Module? - -Previously, `model.py` was in `coldpress/` and only the `coldpress` tool validated configs. The `coldpress-setup` tool would accept invalid YAMLs and fail during cluster application. Now both tools validate upfront using the same shared models. - -## Migration from Legacy - -This replaces the old validation system with modern Pydantic 2.0+ models. The validation is: - -- More comprehensive -- Better error messages -- Type-safe -- Easier to extend -- Compatible with modern Python tooling -- **Shared between both tools** (new!) - -## Adding New Validations - -To add new validation rules, edit `coldpress_common/model.py`: - -1. Add fields to the appropriate Pydantic model -2. Add `@field_validator` or `@model_validator` decorators for custom logic -3. Export the function in `coldpress_common/__init__.py` -4. Update this documentation -5. Add test cases to `test_validation.py` - -Example: - -```python -# In coldpress_common/model.py -class TaskSpec(BaseModel): - name: str - - @field_validator("name") - @classmethod - def validate_name_format(cls, v): - if not v.islower(): - raise ValueError("Task name must be lowercase") - return v -``` diff --git a/docs/quickstart_admin.md b/docs/quickstart_admin.md deleted file mode 100644 index 66ea8a7..0000000 --- a/docs/quickstart_admin.md +++ /dev/null @@ -1,213 +0,0 @@ - -# Coldpress Admin Quickstart - -This guide covers cluster-wide configuration tasks that require admin privileges. These are typically done once by a cluster administrator. - -**Prerequisites:** -- Admin access to Kubernetes/OpenShift cluster (tested on OpenShift 4.21.5, Kubernetes v1.34.4) -- Kueue operator installed (tested with v0.11.6, API v1beta1) -- JobSet operator installed (tested with v1.0.0, API v1alpha2) -- `kubectl` or `oc` CLI installed and configured (tested with oc 4.17.0) -- Coldpress CLI tools installed (see main [README.md](../README.md)) - ---- - -## Step 1: Generate and Apply Cluster Configuration - -**What:** Generate manifests for cluster-wide resources (ResourceFlavors, ClusterQueue). - -**Why:** Sets up the Kueue queueing system for GPU allocation across the cluster. - -**Step 1a: Generate manifests** -```bash -coldpress-setup generate cluster ocp-test-nerc-mghpcc.yaml -``` - -This generates: -- `manifests/cluster-ocp-test-nerc-mghpcc-.yaml` - Kubernetes manifests -- `manifests/label-nodes-ocp-test-nerc-mghpcc-.sh` - Node labeling script - -**Step 1b: Review the generated manifests** -```bash -cat manifests/cluster-ocp-test-nerc-mghpcc-*.yaml -cat manifests/label-nodes-ocp-test-nerc-mghpcc-*.sh -``` - -**Step 1c: Run the labeling script** - -Coldpress uses node labels for scheduling jobs to specific GPU nodes. - -```bash -./manifests/label-nodes-ocp-test-nerc-mghpcc-*.sh -``` - -**Step 1d: Apply the manifest to the cluster** -```bash -oc apply -f manifests/cluster-ocp-test-nerc-mghpcc-*.yaml -``` - -**What gets created:** -- **Node labels** (via labeling script in Step 1c): `coldpress.node: 0, 1, etc.` -- **ResourceFlavors** (via manifest in Step 1d): GPU node pools referencing labeled nodes -- **ClusterQueue** (via manifest in Step 1d): cluster-queue-coldpress - -**Note:** The Kueue and JobSet operators must already be installed on the cluster. This step only creates the Kueue custom resources that use those operators. - -**Verification:** -```bash -# Verify node labels were applied -oc get nodes --show-labels | grep coldpress.node - -# Verify cluster resources were created -oc get clusterqueues -oc get resourceflavors -``` - -**You should see:** -``` -# Node labels -wrk-4 ... coldpress.node=0 ... -wrk-6 ... coldpress.node=1 ... - -# Cluster resources -NAME AGE -cluster-queue-coldpress 5s - -NAME AGE -node0 5s -node1 5s -``` - -**About node labeling:** - -The cluster configuration's `nodes` section specifies which nodes should be labeled with coldpress IDs. Node IDs are automatically assigned based on the order in the YAML file (0, 1, 2, ...). The generated labeling script applies these labels, which are required for coldpress to schedule jobs to specific GPU nodes. - ---- - -## Step 2: Generate and Apply Project Configuration - -**What:** Generate manifests for a project namespace with storage and queueing resources. - -**Why:** Provides isolated workspace for a research group or project team. - -**Step 2a: Generate manifests** -```bash -coldpress-setup generate project coldpress-project.yaml -``` - -This generates a timestamped manifest file (e.g., `manifests/project-coldpress-project-20260413-152928.yaml`). - -**Step 2b: Review the generated manifest** -```bash -cat manifests/project-coldpress-project-*.yaml -``` - -Review the RBAC permissions and resource allocations before applying. - -**Step 2c: Apply the manifest to the cluster** -```bash -oc apply -f manifests/project-coldpress-project-*.yaml -``` - -**This creates:** -- Namespace: `coldpress-project` -- LocalQueue: `coldpress-local-queue-coldpress-project` (connects to ClusterQueue) -- PersistentVolumeClaim: `coldpress-project-storage` (500Gi) -- RBAC: ServiceAccount, Role, RoleBinding for job execution - -**Verification:** -```bash -oc get namespace coldpress-project -oc get pvc -n coldpress-project -oc get localqueues -n coldpress-project -``` - -**You should see:** -``` -NAME STATUS AGE -coldpress-project Active 5s - -NAME STATUS VOLUME CAPACITY ACCESS MODES AGE -coldpress-project-storage Bound pvc-... 500Gi RWX 5s - -NAME CLUSTERQUEUE AGE -coldpress-local-queue-coldpress-project cluster-queue-coldpress 5s -``` - ---- - -## Step 3: Generate and Apply User RBAC - -**What:** Generate manifests to grant an existing cluster user permission to submit jobs to the project namespace. - -**Why:** Allows regular users to create and manage JobSets without admin privileges. - -**Prerequisites:** -- User must already exist in the cluster's authentication system (OpenShift OAuth, LDAP, etc.) -- Project configuration must be applied first (creates the Role that this RoleBinding references) - -**Step 3a: Generate manifests** -```bash -coldpress-setup generate user coldpress-user.yaml -``` - -This generates a timestamped manifest file (e.g., `manifests/user-coldpress-user-20260413-153047.yaml`). - -**User config example** (coldpress-user.yaml in the users/ directory): -```yaml -username: coldpress-user -namespaces: - - coldpress-project -``` - -**Step 3b: Review the generated manifest** -```bash -cat manifests/user-coldpress-user-*.yaml -``` - -**Step 3c: Apply the manifest to the cluster** -```bash -oc apply -f manifests/user-coldpress-user-*.yaml -``` - -**This creates:** -- RoleBinding: `coldpress-user-coldpress-user` in namespace `coldpress-project` -- Binds existing user to existing Role: `coldpress-user-role` (created by project setup) -- Grants permissions: create/manage JobSets, view Jobs/Pods/Services - -**Important:** This does not create a user account. Users must already exist in your cluster's authentication system. - -**Verification:** -```bash -oc get rolebindings -n coldpress-project | grep coldpress-user -``` - -**Expected output:** -``` -coldpress-user-coldpress-user 5s -``` - ---- - -## Summary - -You have now completed the admin setup for Coldpress: - -1. ✓ Applied node labels (required for job scheduling) -2. ✓ Configured cluster-wide Kueue resources (ClusterQueue, ResourceFlavors) -3. ✓ Set up project namespace with storage and queueing (LocalQueue, PVC) -4. ✓ Configured user RBAC for job submission - -**Next steps:** - -Users can now follow the [User Quickstart Guide](quickstart_user.md) to submit and manage AI/HPC workloads. - -**For additional users:** -1. Create a user config file in the `users/` directory (e.g., username.yaml) -2. Generate manifests with `coldpress-setup generate user username.yaml` -3. Apply the manifest with `oc apply -f manifests/user-*.yaml` - -**For additional projects:** -1. Create a project config file in the `projects/` directory (e.g., project-name.yaml) -2. Generate manifests with `coldpress-setup generate project project-name.yaml` -3. Apply the manifest with `oc apply -f manifests/project-*.yaml` diff --git a/docs/quickstart_user.md b/docs/quickstart_user.md deleted file mode 100644 index 460f4e2..0000000 --- a/docs/quickstart_user.md +++ /dev/null @@ -1,421 +0,0 @@ - -# Coldpress User Quickstart - -This guide shows the typical workflow for a regular user submitting and managing an AI/HPC workload. These steps are repeatable for each job. - -**Prerequisites:** -- **Admin must have completed the [Admin Quickstart Guide](quickstart_admin.md) first** (cluster, project, and user configuration) -- Coldpress CLI tools installed (see main [README.md](../README.md)) -- `kubectl` or `oc` CLI installed and configured -- User has access to the target namespace - ---- - -## Step 1: Generate Job Manifests - -**What:** Generate JobSet manifests and helper scripts from a job specification. - -**Why:** Creates all Kubernetes resources needed to run your workload. - -**Command:** -```bash -coldpress generate --config examples/pytorch_ddp_training/config.yaml -``` - -> **Note:** If using uv, activate the venv first with `source .venv/bin/activate`. If using pipx, no activation needed. - -**Input files:** -- `examples/pytorch_ddp_training/config.yaml` - Coldpress job configuration -- `examples/pytorch_ddp_training/job-spec.yaml` - Workload specification -- `examples/pytorch_ddp_training/train.py` - Training script -- `examples/pytorch_ddp_training/model_config.json` - Model configuration - -**This generates** (`output/ddp-training-job/`): -``` -ddp-training-job/ -├── jobset.yaml # JobSet manifest with all tasks -├── metadata.json # Job metadata and node assignments -├── run.sh # Apply JobSet and wait for completion -├── cleanup.sh # Delete JobSet and services -├── monitor.sh # Watch job status -├── logs.sh # Capture and save logs to PVC -├── explore.sh # Interactive shell to browse results -└── cp.sh # Copy results from PVC to local directory -``` - -**You'll see output like:** -``` -Generating JobSet for: ddp-training -Project: coldpress-project -Namespace: coldpress-project -Tasks: 1 - - Task 0 (ddp-training) → Node 0 (GPUs: 2) - -Generated: ddp-training-job/jobset.yaml -Generated: ddp-training-job/metadata.json -Generated: ddp-training-job/run.sh -Generated: ddp-training-job/cleanup.sh -Generated: ddp-training-job/monitor.sh -Generated: ddp-training-job/logs.sh -Generated: ddp-training-job/explore.sh -Generated: ddp-training-job/cp.sh -``` - -**Result:** Job manifests and helper scripts are now ready to use. - ---- - -## Step 2: Inspect Generated Manifests (Optional) - -**What:** Review the generated JobSet manifest before applying. - -**Why:** Understand what resources will be created and verify configuration. - -**File:** `output/ddp-training-job/jobset.yaml` - -**JobSet structure:** -```yaml -apiVersion: jobset.x-k8s.io/v1alpha2 -kind: JobSet -metadata: - name: coldpress-ddp-training - namespace: coldpress-project - labels: - kueue.x-k8s.io/queue-name: coldpress-local-queue-coldpress-project -spec: - replicatedJobs: - - name: mkdir # Job 1: Create results directory and task subdirectories - - name: task-0 # Job 2: PyTorch DDP training (2 GPUs, with discovery init container) -``` - -**Key configuration for task-0 (training job):** -- Image: `pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime` -- Node selector: `coldpress.node: '0'` -- Resources: 2 GPUs, 16Gi memory, 8 CPU cores -- Command: `python -m torch.distributed.run --nproc_per_node=2 train.py` -- Volumes: PVC for results, emptyDir for shared memory -- Dependencies: Waits for mkdir job to complete -- Init container: Runs discovery to capture node hardware before training starts - ---- - -## Step 3: Run the Job - -**What:** Apply the JobSet to the cluster and wait for completion. - -**Why:** Submits your workload to Kueue for scheduling and execution. - -**Commands:** -```bash -cd output/ddp-training-job -./run.sh -``` - -**What happens:** -1. JobSet is created in the cluster -2. Kueue queues the job and waits for resources -3. When GPUs are available, Kueue unsuspends the JobSet -4. Jobs execute in order: mkdir → training (with discovery init container) -5. Script waits for all jobs to complete - -**You'll see:** -``` -Applying JobSet... -jobset.jobset.x-k8s.io/ddp-training created - -Waiting for JobSet to complete... -Job status: Running -Job status: Running -... -Job status: Complete - -JobSet completed successfully! -``` - -**Execution timeline (typical):** -- mkdir: 5-10 seconds (creates base dir + task subdirectories) -- task-0 init container (discovery): 5-10 seconds (runs before main container) -- task-0 main container (training): 2-3 minutes (depends on workload) - -**Result:** Your job is now submitted and will execute when resources are available. - ---- - -## Step 4: Monitor Job Progress - -**What:** Watch the job status in real-time. - -**Why:** Track progress and identify issues quickly. - -**Command:** -```bash -./monitor.sh -``` - -**Output:** -``` -Monitoring JobSet: ddp-training - -NAMESPACE NAME READY AGE -coldpress-project ddp-training-mkdir-0 1/1 15s -coldpress-project ddp-training-task-0-0 0/1 25s - -Pods: -NAME READY STATUS RESTARTS AGE -ddp-training-mkdir-0-0-abc123 0/1 Completed 0 15s -ddp-training-task-0-0-ghi789 0/1 Init:0/1 0 10s # Discovery running -ddp-training-task-0-0-ghi789 1/1 Running 0 25s # Training started -``` - -**Tip:** Press Ctrl+C to exit monitoring. - ---- - -## Step 5: View and Save Logs - -**What:** Capture pod logs and save them to persistent storage. - -**Why:** Preserve training output, metrics, and debugging information. - -**Command:** -```bash -./logs.sh -``` - -**Output:** -``` -Capturing logs for job: ddp-training -Fetching logs from pod: ddp-training-task-0-0-ghi789 - -Logs saved to PVC: - /data/coldpress-project/coldpress_results/ddp-training-9bdbf55a-20260409_072200/logs/ - ├── ddp-training-task-0-0-ghi789.log - └── combined.log - -Log capture complete! -``` - -**What's in the logs:** -- Dataset download progress -- NCCL initialization (GPU communication) -- Training progress (epochs, loss, accuracy) -- Model save confirmation - -**Example log snippet:** -``` -Epoch 10/50 - Loss: 0.234 - Accuracy: 85.2% -Epoch 20/50 - Loss: 0.156 - Accuracy: 90.8% -Epoch 30/50 - Loss: 0.128 - Accuracy: 92.5% -Epoch 40/50 - Loss: 0.115 - Accuracy: 93.8% -Epoch 50/50 - Loss: 0.111 - Accuracy: 94.55% -Saving model to /results/checkpoints/model_weights.pth -Training complete! -``` - -**Result:** Logs are now captured and saved to your PVC. - ---- - -## Step 6: Explore Results - -**What:** Browse the results directory in persistent storage. - -**Why:** Inspect training outputs, model weights, and metrics. - -### Option 1: Using explore.sh (recommended) - -```bash -./explore.sh -``` - -**What happens:** -- Creates a temporary interactive pod -- Mounts the PVC with your results -- Opens a shell at the results directory -- Auto-cleans up the pod when you exit - -**Inside the shell:** -```bash -# You're now in the results directory -ls -lh - -# Output: -# task-0/ -# logs/ - -# Navigate to task-0 directory -cd task-0 -ls -lh - -# Output: -# discovery_user_snapshot.json -# checkpoints/ - -# Check training stats -cat checkpoints/training_stats.json - -# View model file -ls -lh checkpoints/model_weights.pth -# -rw-r--r-- 1 nobody nobody 77M Apr 9 07:24 model_weights.pth - -# Exit the shell -exit -``` - -### Option 2: Quick check with oc - -```bash -oc run check-results --rm -i --restart=Never \ - --image=ubi9/ubi-minimal -n coldpress-project \ - --overrides='{"spec":{"volumes":[{"name":"data","persistentVolumeClaim":{"claimName":"coldpress-project-storage"}}],"containers":[{"name":"check","image":"ubi9/ubi-minimal","command":["ls","-lR","/data/coldpress-project/coldpress_results"],"volumeMounts":[{"name":"data","mountPath":"/data"}]}]}}' -``` - -### Results directory structure - -``` -/data/coldpress-project/coldpress_results/ddp-training-{uid}-{timestamp}/ -├── task-0/ -│ ├── discovery_user_snapshot.json # Hardware/benchmark data (2.7KB) -│ └── checkpoints/ -│ ├── model_weights.pth # Trained model (77MB) -│ └── training_stats.json # Training metrics (303 bytes) -└── logs/ - ├── ddp-training-task-0-0-ghi789.log # Individual pod log (8.7KB) - └── combined.log # Combined logs (8.8KB) -``` - -### Training statistics example - -**File:** `checkpoints/training_stats.json` -```json -{ - "dataset": "mnist", - "epochs": 50, - "batch_size": 128, - "hidden_size": 4096, - "train_test_split": 0.8, - "num_gpus": 2, - "time_seconds": 116.05, - "final_loss": 0.111, - "accuracy": 94.55, - "model_params": 20037642, - "input_dim": 784, - "output_dim": 10 -} -``` - -**Result:** When training completes, you should see results like 94.55% accuracy achieved in ~2 minutes. - ---- - -## Step 7: Copy Results to Local Machine - -**What:** Copy results from the PVC to your local machine. - -**Why:** Download results for local analysis, backup, or sharing. - -**Command:** -```bash -# Copy to default location (./results) -./cp.sh - -# Copy to specific directory -./cp.sh /path/to/destination -``` - -**Output:** -``` -Copying results from PVC to local directory... -Creating temporary pod... -Copying files... -Cleaning up temporary pod... - -===== Copy Complete ===== -Results copied to: ./results -========================== -``` - -**Result:** Results are now available on your local machine. - ---- - -## Step 8: Cleanup Resources - -**What:** Delete the JobSet and associated Kubernetes resources. - -**Why:** Free up cluster resources while preserving results in persistent storage. - -**Command:** -```bash -./cleanup.sh -``` - -**This will delete:** -- JobSet: `ddp-training` -- All Jobs and Pods (cascading delete) -- ConfigMap: `ddp-training-files` (injected files) -- Services (if any were created) - -**This will preserve:** -- All results in PVC (`coldpress-project-storage`) -- Discovery snapshots -- Model weights and checkpoints -- Training logs and statistics - -**Output:** -``` -Cleaning up resources for job: ddp-training -Deleting JobSet... -jobset.jobset.x-k8s.io "ddp-training" deleted -Deleting Services... -No resources found -Deleting ConfigMap... -configmap "ddp-training-files" deleted -Cleanup complete! -``` - -**Verification:** -```bash -oc get jobset,job,pod,configmap -n coldpress-project | grep ddp-training -# Output: No resources found (all cleaned up) - -oc get pvc -n coldpress-project -# Output: coldpress-project-storage still exists with all results intact -``` - -**Result:** Cluster resources are now cleaned up, with all results preserved in your PVC. - ---- - -## Summary - -By following this guide, you have: - -1. ✓ Generated job manifests from specification -2. ✓ Submitted JobSet to cluster -3. ✓ Monitored job progress -4. ✓ Captured and saved logs -5. ✓ Explored results in persistent storage -6. ✓ Copied results to local machine -7. ✓ Cleaned up Kubernetes resources - -**Final results:** -- Dataset: MNIST -- Configuration: 2 GPUs, 50 epochs, batch size 128 -- Performance: 94.55% accuracy in 116 seconds -- Model size: 77MB (20M parameters) - ---- - -## Next Steps - -**For new workloads:** -1. Create a job specification in `examples/your-workload/` -2. Run `coldpress generate --config examples/your-workload/config.yaml` -3. Follow the steps in this guide to run your job - -**Advanced features:** -- See [examples/README.md](../examples/README.md) for more example workloads -- See [VALIDATION.md](VALIDATION.md) for YAML validation system documentation -- See main [README.md](../README.md) for advanced features like node scheduling diff --git a/examples/README.md b/examples/README.md deleted file mode 100644 index cfcdadc..0000000 --- a/examples/README.md +++ /dev/null @@ -1,180 +0,0 @@ - -# Coldpress Examples - -Example workloads with their configuration files. Each example directory contains: -- `config.yaml` - Project, discovery, and output settings -- `job-spec.yaml` - Workload specification - -## Usage - -Point coldpress to a config file: - -```bash -coldpress generate --config examples/pytorch_ddp_training/config.yaml -``` - -Override config values with CLI args: - -```bash -coldpress generate \ - --config examples/pytorch_ddp_training/config.yaml \ - --project different-project \ - --output custom-output -``` - -## Config File Format - -Each workload directory contains: -- `config.yaml` - Project, discovery, and output settings -- `job-spec.yaml` - Workload specification (auto-discovered) - -**config.yaml:** -```yaml -project: coldpress-project # References project config coldpress-project.yaml (stored in projects/ directory) - -# Discovery - runs as init container per task to capture actual node hardware -discovery: user_snapshot # Simple format (all tasks, backward compatible) - -# Or use detailed format to control which tasks run discovery: -# discovery: -# template: user_snapshot -# tasks: all # 'all' or list like [0, 1, 2] - -output: ddp-training-job # Output directory -``` - -**Discovery Configuration:** -- **Simple format**: `discovery: user_snapshot` - Runs discovery for all tasks -- **Detailed format**: Specify which tasks run discovery: - - `tasks: all` - All tasks run discovery (default) - - `tasks: [0, 2]` - Only tasks 0 and 2 run discovery -- **Per-task execution**: Discovery runs as init container, capturing hardware of the actual node where each task executes -- **Results location**: `{base_dir}/task-{N}/discovery_{template}.json` - -**Environment Variables:** -- `COLDPRESS_PROJECT_DIR` - Project configs directory (default: `projects`) -- `COLDPRESS_DISCOVERY_DIR` - Discovery templates directory (default: `discovery`) -- `COLDPRESS_OUTPUT_DIR` - Default output directory (default: `output`) - -## Available Examples - -### PyTorch DDP Training - -Distributed PyTorch training with 4 GPUs on MNIST dataset. - -**Run:** -```bash -coldpress generate --config examples/pytorch_ddp_training/config.yaml -cd ddp-training-job/ -./run.sh -``` - -**Job spec (generates JobSet named `coldpress-ddp-training`):** -```yaml -name: ddp-training - -containers: - - name: training - image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime - command: ["python", "-m", "torch.distributed.run"] - args: - - --nproc_per_node=2 - - --nnodes=1 - - train.py - - --dataset=mnist - - --train-test-split=0.8 - - --epochs=50 - - --batch-size=128 - - --hidden-size=4096 - - --lr=0.01 - - --output-dir=/results/checkpoints - resources: - requests: - nvidia.com/gpu: "2" - memory: "16Gi" - cpu: "8" - limits: - nvidia.com/gpu: "2" - memory: "16Gi" - -volumes: - - name: results - mount: /results - - name: dshm - type: emptyDir - medium: Memory - sizeLimit: 16Gi - mount: /dev/shm -``` - -### vLLM + GuideLLM Benchmark - -vLLM inference server + GuideLLM benchmark client for model serving performance testing. - -**Run:** -```bash -coldpress generate --config examples/vllm_guidellm_benchmark/config.yaml -cd vllm-benchmark-job/ -./run.sh -``` - -**Template:** - -Multi-task workflow: vLLM inference server + GuideLLM benchmark client - -```yaml -name: vllm-guidellm - -tasks: - - name: inference-server - blocking: endpoint - health_check: http://127.0.0.1:8000/health - containers: - - name: server - image: quay.io/vllm/vllm:latest - command: ["python", "-m", "vllm.entrypoints.openai.api_server"] - args: - - --model=ibm-granite/granite-3.3-8b-instruct - - --port=8000 - resources: - requests: - nvidia.com/gpu: "1" - volumes: - - name: results - mount: /results - - - name: benchmark-client - blocking: completion - containers: - - name: client - image: ghcr.io/vllm-project/guidellm:nightly - command: ["guidellm"] - args: - - --target=http://inference-server:8000 - - --duration=30s - volumes: - - name: results - mount: /results -``` - -## Creating Your Own - -Create a directory with your workload: - -``` -my-workflow/ - config.yaml # Settings - job-spec.yaml # Workload specification (auto-discovered) -``` - -**config.yaml:** -```yaml -project: your-project -discovery: user_snapshot -output: my-job -``` - -Then generate: -```bash -coldpress generate --config my-workflow/config.yaml -``` diff --git a/examples/pytorch_ddp_training/config.yaml b/examples/pytorch_ddp_training/config.yaml deleted file mode 100644 index 254ede5..0000000 --- a/examples/pytorch_ddp_training/config.yaml +++ /dev/null @@ -1,9 +0,0 @@ -# Generated by: Claude Sonnet 4.5 -project: coldpress-project -discovery: user_snapshot -output: ddp-training-job - -# Files to mount into container (creates ConfigMap) -files: - - train.py - - model_config.json diff --git a/examples/pytorch_ddp_training/intent_jobset.yaml b/examples/pytorch_ddp_training/intent_jobset.yaml new file mode 100644 index 0000000..92fcaec --- /dev/null +++ b/examples/pytorch_ddp_training/intent_jobset.yaml @@ -0,0 +1,33 @@ +project: coldpress-project +output: ddp-training-job +target: jobset + +files: + - train.py + +discovery: + template: user_snapshot + tasks: all + +tasks: + - name: ddp-training + replicas: 2 + + args: + + # Positioned insertion: new args with insert_after + nnodes: + value: "${REPLICAS}" + insert_after: nproc_per_node + + node_rank: + value: "${INDEX}" + insert_after: nnodes + + master_addr: + value: "${REPLICA_ddp-training_0}" + insert_after: node_rank + + master_port: + value: "29500" + insert_after: master_addr diff --git a/examples/pytorch_ddp_training/intent_kubeflow.yaml b/examples/pytorch_ddp_training/intent_kubeflow.yaml new file mode 100644 index 0000000..a574856 --- /dev/null +++ b/examples/pytorch_ddp_training/intent_kubeflow.yaml @@ -0,0 +1,21 @@ +project: coldpress-project +output: ddp-training-job +target: kubeflow + +files: + - train.py + +discovery: + template: user_snapshot + tasks: all + +tasks: + - name: ddp-training + replicas: 2 + + args: + # Positioned insertion: nnodes after nproc_per_node + # PyTorchJob auto-injects MASTER_ADDR, MASTER_PORT, RANK via env vars + nnodes: + value: "${REPLICAS}" + insert_after: nproc_per_node diff --git a/examples/pytorch_ddp_training/job-spec.yaml b/examples/pytorch_ddp_training/job-spec.yaml index 0f198f9..9864052 100644 --- a/examples/pytorch_ddp_training/job-spec.yaml +++ b/examples/pytorch_ddp_training/job-spec.yaml @@ -1,41 +1,61 @@ -# Generated by: Claude Sonnet 4.5 -name: ddp-training - -tolerate_all: true - -containers: - - name: training - image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime - workingDir: /workspace - command: ["python", "-m", "torch.distributed.run"] - args: - - --nproc_per_node=2 - - --nnodes=1 - - train.py - - --dataset=mnist - - --train-test-split=0.8 - - --epochs=50 - - --batch-size=128 - - --hidden-size=4096 - - --lr=0.01 - - --output-dir=/results/checkpoints - resources: - requests: - nvidia.com/gpu: "2" - memory: "16Gi" - cpu: "8" - limits: - nvidia.com/gpu: "2" - memory: "16Gi" - env: - - name: NCCL_DEBUG - value: "INFO" - -volumes: - - name: results - mount: /results - - name: dshm - type: emptyDir - medium: Memory - sizeLimit: 16Gi - mount: /dev/shm +apiVersion: batch/v1 +kind: Job +metadata: + name: ddp-training + namespace: coldpress-project +spec: + template: + metadata: + labels: + app: ddp-training + spec: + restartPolicy: Never + tolerations: + - operator: Exists + containers: + - name: training + image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime + workingDir: /workspace + command: + - python + - -m + - torch.distributed.run + - --nproc_per_node=2 + - train.py + - --dataset=mnist + - --train-test-split=0.8 + - --epochs=50 + - --batch-size=128 + - --hidden-size=4096 + - --lr=0.01 + - --output-dir=/results/checkpoints + resources: + requests: + nvidia.com/gpu: "2" + memory: "16Gi" + cpu: "8" + limits: + nvidia.com/gpu: "2" + memory: "16Gi" + env: + - name: NCCL_DEBUG + value: "INFO" + volumeMounts: + - name: results + mountPath: /results + - name: dshm + mountPath: /dev/shm + - name: training-script + mountPath: /workspace/train.py + subPath: train.py + volumes: + - name: results + persistentVolumeClaim: + claimName: coldpress-project-storage + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 16Gi + - name: training-script + configMap: + name: ddp-training-files diff --git a/examples/pytorch_ddp_training/model_config.json b/examples/pytorch_ddp_training/model_config.json deleted file mode 100644 index 6dbe313..0000000 --- a/examples/pytorch_ddp_training/model_config.json +++ /dev/null @@ -1,8 +0,0 @@ -{ - "model": "SimpleNet", - "hidden_size": 4096, - "dataset": "mnist", - "epochs": 50, - "batch_size": 128, - "learning_rate": 0.01 -} diff --git a/examples/pytorch_ddp_training/train.py b/examples/pytorch_ddp_training/train.py index e4c92c3..e325b23 100644 --- a/examples/pytorch_ddp_training/train.py +++ b/examples/pytorch_ddp_training/train.py @@ -50,25 +50,35 @@ def main(): dist.init_process_group("nccl") rank = dist.get_rank() - # Dataset loading (synchronized to avoid race condition) + # Dataset loading + # Use local /tmp for dataset downloads (each pod downloads its own copy) + # /tmp is pod-local, so only local_rank 0 within each pod should download + dataset_root = "/tmp/data" + transform = transforms.Compose( [transforms.ToTensor(), transforms.Lambda(lambda x: torch.flatten(x))] ) - if rank == 0: + # Only local_rank 0 downloads to avoid race conditions within the pod + if local_rank == 0: if args.dataset.lower() == "mnist": torchvision.datasets.MNIST( - root="/tmp/data", train=True, download=True, transform=transform + root=dataset_root, train=True, download=True, transform=transform ) else: raise ValueError("Unsupported dataset") - dist.barrier() + # Wait for local_rank 0 to finish downloading + if dist.is_initialized(): + dist.barrier() + # All ranks load the dataset if args.dataset.lower() == "mnist": full_dataset = torchvision.datasets.MNIST( - root="/tmp/data", train=True, download=False, transform=transform + root=dataset_root, train=True, download=False, transform=transform ) + else: + raise ValueError("Unsupported dataset") input_dim = full_dataset[0][0].numel() output_dim = len(full_dataset.classes) if hasattr(full_dataset, "classes") else 10 diff --git a/examples/pytorch_ddp_training_node1/README.md b/examples/pytorch_ddp_training_node1/README.md deleted file mode 100644 index d3bb985..0000000 --- a/examples/pytorch_ddp_training_node1/README.md +++ /dev/null @@ -1,135 +0,0 @@ -# PyTorch DDP Training - Explicit Node Assignment - -This example demonstrates running a PyTorch Distributed Data Parallel (DDP) training job with **explicit node assignment** specified in the config file. - -## Difference from Base Example - -**Base example (`pytorch_ddp_training`):** -- No `nodes` field in config.yaml -- Kubernetes scheduler picks any node with `coldpress.node` label -- Uses `nodeAffinity` with `Exists` operator - -**This example (`pytorch_ddp_training_node1`):** -- Has `nodes: [1]` in config.yaml -- Task is pinned to node 1 (label: `coldpress.node=1`) -- Uses `nodeSelector` with specific node ID - -## Config File - -```yaml -# config.yaml -project: coldpress-project - -# Per-task discovery - runs on the actual node where task executes -discovery: - template: user_snapshot - tasks: all # Run discovery for all tasks - -output: ddp-training-job-node1 - -files: - - train.py - - model_config.json - -# Explicit node assignment - pin task to node 1 -nodes: - - 1 -``` - -## Per-Task Discovery - -This example uses **per-task discovery**, which runs hardware discovery as an init container for each task. This ensures the discovery snapshot reflects the **actual hardware** where the task executes, not a different node. - -**Discovery configuration:** -- `template: user_snapshot` - Discovery template to use -- `tasks: all` - Run discovery for all tasks (can also be `[0, 1, 2]` for specific tasks) - -**Results structure:** -``` -/data/coldpress-project/coldpress_results/ddp-training-job-{uid}-{timestamp}/ -├── task-0/ -│ └── discovery_user_snapshot.json # Hardware info for task 0 -└── logs/ - └── task-0.log -``` - -**Backward compatibility:** Simple string format still works: -```yaml -discovery: user_snapshot # Equivalent to {template: user_snapshot, tasks: all} -``` - -## Usage - -```bash -# Generate manifests (node assignment is in config.yaml) -coldpress generate --config examples/pytorch_ddp_training_node1/config.yaml - -# Output directory -cd output/ddp-training-job-node1/ - -# Apply to cluster -./run.sh - -# Monitor -./monitor.sh -``` - -## Generated Manifest Difference - -**With `--node 1`:** -```yaml -spec: - nodeSelector: - coldpress.node: "1" -``` - -**Without `--node` (base example):** -```yaml -spec: - affinity: - nodeAffinity: - requiredDuringSchedulingIgnoredDuringExecution: - nodeSelectorTerms: - - matchExpressions: - - key: coldpress.node - operator: Exists -``` - -## When to Use Explicit Node Assignment - -Use `--node` when you need: -- **Specific hardware**: Task requires GPU model only available on certain nodes -- **Reproducibility**: Pin experiments to same hardware for consistent results -- **Resource isolation**: Dedicate specific nodes to specific teams/projects -- **Multi-task workflows**: Each task has different hardware requirements - -Otherwise, let the Kubernetes scheduler handle allocation for better cluster utilization. - -## Multi-Task Example - -For workflows with multiple tasks on different nodes, you can specify in config.yaml: - -```yaml -# config.yaml for multi-task workflow -project: coldpress-project -output: multi-task-job - -nodes: - - 0 # Task 0 → node 0 - - 1 # Task 1 → node 1 -``` - -Or override via CLI: -```bash -coldpress generate --config job.yaml --node 0 --node 1 -``` - -**Priority:** CLI `--node` flag > `nodes` in config.yaml > Kubernetes scheduler - -## Files - -Same as base example: -- `config.yaml` - Coldpress configuration -- `job-spec.yaml` - Task specification -- `train.py` - PyTorch training script -- `model_config.json` - Model hyperparameters diff --git a/examples/pytorch_ddp_training_node1/config.yaml b/examples/pytorch_ddp_training_node1/config.yaml deleted file mode 100644 index 49819ea..0000000 --- a/examples/pytorch_ddp_training_node1/config.yaml +++ /dev/null @@ -1,18 +0,0 @@ -# Generated by: Claude Sonnet 4.5 -project: coldpress-project - -# Per-task discovery - runs on the actual node where task executes -discovery: - template: user_snapshot - tasks: all # Run discovery for all tasks (can also be [0, 1, 2] for specific tasks) - -output: ddp-training-job-node1 - -# Files to mount into container (creates ConfigMap) -files: - - train.py - - model_config.json - -# Explicit node assignment - pin task to node 1 -nodes: - - 1 diff --git a/examples/pytorch_ddp_training_node1/job-spec.yaml b/examples/pytorch_ddp_training_node1/job-spec.yaml deleted file mode 100644 index 0f198f9..0000000 --- a/examples/pytorch_ddp_training_node1/job-spec.yaml +++ /dev/null @@ -1,41 +0,0 @@ -# Generated by: Claude Sonnet 4.5 -name: ddp-training - -tolerate_all: true - -containers: - - name: training - image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime - workingDir: /workspace - command: ["python", "-m", "torch.distributed.run"] - args: - - --nproc_per_node=2 - - --nnodes=1 - - train.py - - --dataset=mnist - - --train-test-split=0.8 - - --epochs=50 - - --batch-size=128 - - --hidden-size=4096 - - --lr=0.01 - - --output-dir=/results/checkpoints - resources: - requests: - nvidia.com/gpu: "2" - memory: "16Gi" - cpu: "8" - limits: - nvidia.com/gpu: "2" - memory: "16Gi" - env: - - name: NCCL_DEBUG - value: "INFO" - -volumes: - - name: results - mount: /results - - name: dshm - type: emptyDir - medium: Memory - sizeLimit: 16Gi - mount: /dev/shm diff --git a/examples/pytorch_ddp_training_node1/model_config.json b/examples/pytorch_ddp_training_node1/model_config.json deleted file mode 100644 index 6dbe313..0000000 --- a/examples/pytorch_ddp_training_node1/model_config.json +++ /dev/null @@ -1,8 +0,0 @@ -{ - "model": "SimpleNet", - "hidden_size": 4096, - "dataset": "mnist", - "epochs": 50, - "batch_size": 128, - "learning_rate": 0.01 -} diff --git a/examples/pytorch_ddp_training_node1/train.py b/examples/pytorch_ddp_training_node1/train.py deleted file mode 100644 index e4c92c3..0000000 --- a/examples/pytorch_ddp_training_node1/train.py +++ /dev/null @@ -1,170 +0,0 @@ -# Assisted by: Claude Sonnet 4.5 -"""PyTorch DDP training example - based on demo.""" - -import os -import json -import time -import argparse -import torch -import torch.nn as nn -import torch.optim as optim -import torch.distributed as dist -from torch.nn.parallel import DistributedDataParallel as DDP -from torch.utils.data import DataLoader, random_split -from torch.utils.data.distributed import DistributedSampler -import torchvision -import torchvision.transforms as transforms - - -class DynamicMLP(nn.Module): - """Simple MLP for classification.""" - - def __init__(self, input_dim, hidden_size, output_dim): - super().__init__() - self.layers = nn.Sequential( - nn.Linear(input_dim, hidden_size), - nn.ReLU(), - nn.Linear(hidden_size, hidden_size), - nn.ReLU(), - nn.Linear(hidden_size, output_dim), - ) - - def forward(self, x): - return self.layers(x) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--dataset", type=str, default="mnist") - parser.add_argument("--train-test-split", type=float, default=0.8) - parser.add_argument("--epochs", type=int, default=10) - parser.add_argument("--batch-size", type=int, default=64) - parser.add_argument("--hidden-size", type=int, default=128) - parser.add_argument("--lr", type=float, default=0.01) - parser.add_argument("--output-dir", type=str, default="/results") - args = parser.parse_args() - - # Initialize DDP - local_rank = int(os.environ["LOCAL_RANK"]) - torch.cuda.set_device(local_rank) - dist.init_process_group("nccl") - rank = dist.get_rank() - - # Dataset loading (synchronized to avoid race condition) - transform = transforms.Compose( - [transforms.ToTensor(), transforms.Lambda(lambda x: torch.flatten(x))] - ) - - if rank == 0: - if args.dataset.lower() == "mnist": - torchvision.datasets.MNIST( - root="/tmp/data", train=True, download=True, transform=transform - ) - else: - raise ValueError("Unsupported dataset") - - dist.barrier() - - if args.dataset.lower() == "mnist": - full_dataset = torchvision.datasets.MNIST( - root="/tmp/data", train=True, download=False, transform=transform - ) - - input_dim = full_dataset[0][0].numel() - output_dim = len(full_dataset.classes) if hasattr(full_dataset, "classes") else 10 - - # Train-test split - train_size = int(args.train_test_split * len(full_dataset)) - test_size = len(full_dataset) - train_size - generator = torch.Generator().manual_seed(42) - train_dataset, test_dataset = random_split( - full_dataset, [train_size, test_size], generator=generator - ) - - train_sampler = DistributedSampler(train_dataset) - train_loader = DataLoader( - train_dataset, batch_size=args.batch_size, sampler=train_sampler - ) - test_loader = DataLoader(test_dataset, batch_size=args.batch_size) - - # Model - model = DynamicMLP(input_dim, args.hidden_size, output_dim).cuda(local_rank) - model = DDP(model, device_ids=[local_rank]) - optimizer = optim.SGD(model.parameters(), lr=args.lr) - criterion = nn.CrossEntropyLoss() - - # Training loop - print only every 10 epochs - start = time.time() - accuracy = 0.0 - - for epoch in range(args.epochs): - model.train() - train_sampler.set_epoch(epoch) - running_loss = 0.0 - - for data, target in train_loader: - data, target = data.cuda(local_rank), target.cuda(local_rank) - optimizer.zero_grad() - output = model(data) - loss = criterion(output, target) - loss.backward() - optimizer.step() - running_loss += loss.item() - - # Evaluate and print every 10 epochs (and final epoch) - if (epoch + 1) % 10 == 0 or (epoch + 1) == args.epochs: - model.eval() - correct = 0 - total = 0 - with torch.no_grad(): - for data, target in test_loader: - data, target = data.cuda(local_rank), target.cuda(local_rank) - outputs = model(data) - _, predicted = torch.max(outputs.data, 1) - total += target.size(0) - correct += (predicted == target).sum().item() - - if rank == 0: - avg_loss = running_loss / len(train_loader) - accuracy = 100 * correct / total - print( - f"Epoch [{epoch + 1}/{args.epochs}] - Train Loss: {avg_loss:.4f} - Test Accuracy: {accuracy:.2f}%", - flush=True, - ) - - end = time.time() - - # Save results (rank 0 only) - if rank == 0: - num_params = sum(p.numel() for p in model.parameters()) - - results = { - "dataset": args.dataset, - "epochs": args.epochs, - "batch_size": args.batch_size, - "hidden_size": args.hidden_size, - "train_test_split": args.train_test_split, - "num_gpus": dist.get_world_size(), - "time_seconds": end - start, - "final_loss": float(loss), - "accuracy": accuracy, - "model_params": num_params, - "input_dim": input_dim, - "output_dim": output_dim, - } - - os.makedirs(args.output_dir, exist_ok=True) - - # Save stats - with open(f"{args.output_dir}/training_stats.json", "w") as f: - json.dump(results, f, indent=2) - - # Save model weights - torch.save(model.module.state_dict(), f"{args.output_dir}/model_weights.pth") - print(f"Saved model weights to {args.output_dir}/model_weights.pth", flush=True) - - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/examples/pytorch_ray_training/intent_kuberay.yaml b/examples/pytorch_ray_training/intent_kuberay.yaml new file mode 100644 index 0000000..a357be4 --- /dev/null +++ b/examples/pytorch_ray_training/intent_kuberay.yaml @@ -0,0 +1,20 @@ +project: coldpress-project +output: ray-training-job +target: kuberay + +files: + - train.py + +discovery: + template: user_snapshot + tasks: all + +tasks: + - name: ray-training + replicas: 2 # Coldpress transforms single-node (2 GPU) into 2-pod cluster (1 head + 1 worker = 4 GPUs total) + + args: + # Scale Ray Train workers to match total GPUs (2 pods × 2 GPUs = 4 workers) + num-workers: "4" # Total GPUs across all pods + gpus-per-worker: "1" # 1 GPU per worker for proper device handling + cpus-per-worker: "2" # 2 CPUs per worker diff --git a/examples/pytorch_ray_training/job-spec.yaml b/examples/pytorch_ray_training/job-spec.yaml new file mode 100644 index 0000000..fdd2967 --- /dev/null +++ b/examples/pytorch_ray_training/job-spec.yaml @@ -0,0 +1,61 @@ +apiVersion: batch/v1 +kind: Job +metadata: + name: ray-training + namespace: coldpress-project +spec: + template: + metadata: + labels: + app: ray-training + spec: + restartPolicy: Never + tolerations: + - operator: Exists + containers: + - name: training + image: rayproject/ray-ml:2.9.0-py310-gpu + workingDir: /workspace + command: + - python + - train.py + - --dataset=mnist + - --train-test-split=0.8 + - --epochs=50 + - --batch-size=128 + - --hidden-size=4096 + - --lr=0.01 + - --output-dir=/results/checkpoints + - --num-workers=1 + - --gpus-per-worker=2 + - --cpus-per-worker=4 + resources: + requests: + nvidia.com/gpu: "2" + memory: "16Gi" + cpu: "8" + limits: + nvidia.com/gpu: "2" + memory: "16Gi" + env: + - name: NCCL_DEBUG + value: "INFO" + volumeMounts: + - name: results + mountPath: /results + - name: dshm + mountPath: /dev/shm + - name: training-script + mountPath: /workspace/train.py + subPath: train.py + volumes: + - name: results + persistentVolumeClaim: + claimName: coldpress-project-storage + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 16Gi + - name: training-script + configMap: + name: ray-training-files diff --git a/examples/pytorch_ray_training/train.py b/examples/pytorch_ray_training/train.py new file mode 100644 index 0000000..67801de --- /dev/null +++ b/examples/pytorch_ray_training/train.py @@ -0,0 +1,256 @@ +"""Single-node PyTorch MNIST training using Ray Train - Coldpress will distribute this.""" + +import os +import json +import time +import argparse +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader, random_split +import torchvision +import torchvision.transforms as transforms + +import ray +from ray import train +from ray.train import ScalingConfig, RunConfig, CheckpointConfig +from ray.train.torch import TorchTrainer + + +class DynamicMLP(nn.Module): + """Simple MLP for classification.""" + + def __init__(self, input_dim, hidden_size, output_dim): + super().__init__() + self.layers = nn.Sequential( + nn.Linear(input_dim, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, output_dim), + ) + + def forward(self, x): + return self.layers(x) + + +def train_func(config): + """Training function for Ray Train - runs on each worker. + + Ray Train automatically handles: + - Device placement + - Model DDP wrapping + - Data distribution + """ + # Get hyperparameters + dataset = config.get("dataset", "mnist") + train_test_split = config.get("train_test_split", 0.8) + epochs = config.get("epochs", 10) + batch_size = config.get("batch_size", 64) + hidden_size = config.get("hidden_size", 128) + lr = config.get("lr", 0.01) + output_dir = config.get("output_dir", "/results") + + # Ray Train context + rank = train.get_context().get_world_rank() + world_size = train.get_context().get_world_size() + + # Dataset loading - shared storage + dataset_root = "/results/datasets" + transform = transforms.Compose( + [transforms.ToTensor(), transforms.Lambda(lambda x: torch.flatten(x))] + ) + + # Only rank 0 downloads + if rank == 0: + if dataset.lower() == "mnist": + torchvision.datasets.MNIST( + root=dataset_root, train=True, download=True, transform=transform + ) + else: + raise ValueError("Unsupported dataset") + + # Wait for download + torch.distributed.barrier() + + # All ranks load + if dataset.lower() == "mnist": + full_dataset = torchvision.datasets.MNIST( + root=dataset_root, train=True, download=False, transform=transform + ) + else: + raise ValueError("Unsupported dataset") + + input_dim = full_dataset[0][0].numel() + output_dim = len(full_dataset.classes) if hasattr(full_dataset, "classes") else 10 + + # Train-test split + train_size = int(train_test_split * len(full_dataset)) + test_size = len(full_dataset) - train_size + generator = torch.Generator().manual_seed(42) + train_dataset, test_dataset = random_split( + full_dataset, [train_size, test_size], generator=generator + ) + + # Data loaders + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size) + + # Ray Train handles device placement automatically + train_loader = train.torch.prepare_data_loader(train_loader) + test_loader = train.torch.prepare_data_loader(test_loader) + + # Model - Ray Train wraps in DDP + model = DynamicMLP(input_dim, hidden_size, output_dim) + model = train.torch.prepare_model(model) + + optimizer = optim.SGD(model.parameters(), lr=lr) + criterion = nn.CrossEntropyLoss() + + # Training loop + start = time.time() + accuracy = 0.0 + + for epoch in range(epochs): + model.train() + running_loss = 0.0 + + for data, target in train_loader: + # Ray Train handles device placement + optimizer.zero_grad() + output = model(data) + loss = criterion(output, target) + loss.backward() + optimizer.step() + running_loss += loss.item() + + # Evaluate every 10 epochs + if (epoch + 1) % 10 == 0 or (epoch + 1) == epochs: + model.eval() + correct = 0 + total = 0 + with torch.no_grad(): + for data, target in test_loader: + outputs = model(data) + _, predicted = torch.max(outputs.data, 1) + total += target.size(0) + correct += (predicted == target).sum().item() + + avg_loss = running_loss / len(train_loader) + accuracy = 100 * correct / total + + if rank == 0: + print( + f"Epoch [{epoch + 1}/{epochs}] - Train Loss: {avg_loss:.4f} - Test Accuracy: {accuracy:.2f}%", + flush=True, + ) + + # Report to Ray Train + train.report( + { + "epoch": epoch + 1, + "loss": avg_loss, + "accuracy": accuracy, + } + ) + + end = time.time() + + # Save results (rank 0 only) + if rank == 0: + num_params = sum(p.numel() for p in model.parameters()) + + results = { + "dataset": dataset, + "epochs": epochs, + "batch_size": batch_size, + "hidden_size": hidden_size, + "train_test_split": train_test_split, + "num_gpus": world_size, + "time_seconds": end - start, + "final_loss": float(loss), + "accuracy": accuracy, + "model_params": num_params, + "input_dim": input_dim, + "output_dim": output_dim, + "framework": "ray_train", + } + + os.makedirs(output_dir, exist_ok=True) + + with open(f"{output_dir}/training_stats.json", "w") as f: + json.dump(results, f, indent=2) + + # Save model + model_to_save = model.module if hasattr(model, "module") else model + torch.save(model_to_save.state_dict(), f"{output_dir}/model_weights.pth") + print(f"Saved model weights to {output_dir}/model_weights.pth", flush=True) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", type=str, default="mnist") + parser.add_argument("--train-test-split", type=float, default=0.8) + parser.add_argument("--epochs", type=int, default=10) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--hidden-size", type=int, default=128) + parser.add_argument("--lr", type=float, default=0.01) + parser.add_argument("--output-dir", type=str, default="/results") + # Ray Train scaling config - passed from bash wrapper + parser.add_argument("--num-workers", type=int, default=1) + parser.add_argument("--gpus-per-worker", type=int, default=2) + parser.add_argument("--cpus-per-worker", type=int, default=4) + args = parser.parse_args() + + # Initialize Ray (connects to local Ray or existing cluster) + ray.init() + + # Configure for training + # Single-node default: 1 worker with 2 GPUs + # Distributed: bash wrapper reads env vars and passes as args + scaling_config = ScalingConfig( + num_workers=args.num_workers, + use_gpu=True, + resources_per_worker={"GPU": args.gpus_per_worker, "CPU": args.cpus_per_worker}, + ) + + checkpoint_config = CheckpointConfig( + num_to_keep=1, + checkpoint_score_attribute="accuracy", + checkpoint_score_order="max", + ) + + run_config = RunConfig( + name="mnist-training", + storage_path=args.output_dir, + checkpoint_config=checkpoint_config, + ) + + # Create trainer + trainer = TorchTrainer( + train_func, + train_loop_config={ + "dataset": args.dataset, + "train_test_split": args.train_test_split, + "epochs": args.epochs, + "batch_size": args.batch_size, + "hidden_size": args.hidden_size, + "lr": args.lr, + "output_dir": args.output_dir, + }, + scaling_config=scaling_config, + run_config=run_config, + ) + + # Run training + result = trainer.fit() + + print("\nTraining completed!") + print(f"Best checkpoint: {result.checkpoint}") + print(f"Final metrics: {result.metrics}") + + ray.shutdown() + + +if __name__ == "__main__": + main() diff --git a/examples/vllm_guidellm_benchmark/config.yaml b/examples/vllm_guidellm_benchmark/config.yaml deleted file mode 100644 index c2e1492..0000000 --- a/examples/vllm_guidellm_benchmark/config.yaml +++ /dev/null @@ -1,10 +0,0 @@ -# Generated by: Claude Sonnet 4.5 -project: coldpress-project - -# Per-task discovery - only run on task 0 (inference server with GPU) -# Task 1 (benchmark client) doesn't need hardware discovery -discovery: - template: user_snapshot - tasks: [0] # Only task 0 runs discovery - -output: vllm-benchmark-job diff --git a/examples/vllm_guidellm_benchmark/guidellm-client.yaml b/examples/vllm_guidellm_benchmark/guidellm-client.yaml new file mode 100644 index 0000000..2b40e3a --- /dev/null +++ b/examples/vllm_guidellm_benchmark/guidellm-client.yaml @@ -0,0 +1,52 @@ +apiVersion: batch/v1 +kind: Job +metadata: + name: guidellm-benchmark + namespace: coldpress-project +spec: + template: + metadata: + labels: + app: guidellm-client + spec: + restartPolicy: Never + tolerations: + - operator: Exists + containers: + - name: client + image: ghcr.io/vllm-project/guidellm:nightly + command: ["bash", "-c"] + args: + - | + guidellm benchmark run \ + --target "${GUIDELLM_TARGET}" \ + --output-dir "${GUIDELLM_OUTPUT_DIR}" \ + --outputs "${GUIDELLM_OUTPUTS}" \ + --max-seconds ${GUIDELLM_MAX_SECONDS} \ + --rate-type ${GUIDELLM_RATE_TYPE} \ + --rate ${GUIDELLM_RATE} \ + --data "${GUIDELLM_DATA}" + env: + - name: HOME + value: /tmp + - name: GUIDELLM_TARGET + value: "http://vllm-server:8000" + - name: GUIDELLM_OUTPUT_DIR + value: "/results" + - name: GUIDELLM_OUTPUTS + value: "json" + - name: GUIDELLM_MAX_SECONDS + value: "30" + - name: GUIDELLM_RATE_TYPE + value: "throughput" + - name: GUIDELLM_RATE + value: "1" + - name: GUIDELLM_DATA + value: "prompt_tokens=256,output_tokens=128" + volumeMounts: + - name: results + mountPath: /results + volumes: + - name: results + persistentVolumeClaim: + claimName: coldpress-project-storage diff --git a/examples/vllm_guidellm_benchmark/intent_jobset.yaml b/examples/vllm_guidellm_benchmark/intent_jobset.yaml new file mode 100644 index 0000000..b8c6da2 --- /dev/null +++ b/examples/vllm_guidellm_benchmark/intent_jobset.yaml @@ -0,0 +1,19 @@ +project: coldpress-project +output: vllm-benchmark-job +target: jobset + +discovery: + template: user_snapshot + tasks: all + +tasks: + - name: inference-server + replicas: 1 + + - name: benchmark-client + replicas: 1 + depends_on: + task: inference-server + wait_for: ready + args: + target: "http://${REPLICA_inference-server_0}:8000" diff --git a/examples/vllm_guidellm_benchmark/intent_kserve.yaml b/examples/vllm_guidellm_benchmark/intent_kserve.yaml new file mode 100644 index 0000000..7efc1d3 --- /dev/null +++ b/examples/vllm_guidellm_benchmark/intent_kserve.yaml @@ -0,0 +1,14 @@ +project: coldpress-project +output: vllm-kserve-inference +target: kserve + +discovery: + template: user_snapshot + tasks: all + +tasks: + - name: inference-server + replicas: 1 + env: + # KServe-specific settings could go here + # Environment variables from job-spec.yaml will be preserved diff --git a/examples/vllm_guidellm_benchmark/job-spec.yaml b/examples/vllm_guidellm_benchmark/job-spec.yaml index e85ae0f..dbca378 100644 --- a/examples/vllm_guidellm_benchmark/job-spec.yaml +++ b/examples/vllm_guidellm_benchmark/job-spec.yaml @@ -1,59 +1,95 @@ -# Generated by: Claude Sonnet 4.5 -name: inference-server -tolerate_all: true +# Vanilla Kubernetes manifests +# PROBLEM: These Jobs will race - benchmark-client starts immediately +# and fails because inference-server isn't ready yet. +# Coldpress solves this with JobSet dependencies (dependsOn + status: Ready) -containers: - - name: server - image: nvcr.io/nvidia/vllm:26.03-py3 - command: ["python", "-m", "vllm.entrypoints.openai.api_server"] - args: - - --model=ibm-granite/granite-3.3-8b-instruct - - --port=8000 - - --max-model-len=10000 - - --gpu-memory-utilization=0.6 - env: - - name: HOME - value: /tmp - ports: - - containerPort: 8000 - name: http - readinessProbe: - httpGet: - path: /health - port: 8000 - initialDelaySeconds: 30 - periodSeconds: 10 - resources: - requests: - nvidia.com/gpu: "1" - limits: - nvidia.com/gpu: "1" - -volumes: - - name: results - mount: /results --- -name: benchmark-client -tolerate_all: true - -containers: - - name: client - image: ghcr.io/vllm-project/guidellm:nightly - command: ["bash", "-c"] - args: - - >- - guidellm benchmark run - --target "http://inference-server:8000" - --output-dir "/results" - --outputs "json" - --max-seconds 30 - --rate-type throughput - --rate 1 - --data "prompt_tokens=256,output_tokens=128" - env: - - name: HOME - value: /tmp +apiVersion: batch/v1 +kind: Job +metadata: + name: inference-server + namespace: coldpress-project +spec: + template: + metadata: + labels: + app: inference-server + spec: + restartPolicy: Never + tolerations: + - operator: Exists + containers: + - name: server + image: nvcr.io/nvidia/vllm:26.03-py3 + command: + - python + - -m + - vllm.entrypoints.openai.api_server + - --model=ibm-granite/granite-3.3-8b-instruct + - --port=8000 + - --max-model-len=10000 + - --gpu-memory-utilization=0.6 + env: + - name: HOME + value: /tmp + ports: + - containerPort: 8000 + name: http + readinessProbe: + httpGet: + path: /health + port: 8000 + initialDelaySeconds: 30 + periodSeconds: 10 + resources: + requests: + nvidia.com/gpu: "1" + limits: + nvidia.com/gpu: "1" + volumeMounts: + - name: results + mountPath: /results + volumes: + - name: results + persistentVolumeClaim: + claimName: coldpress-project-storage -volumes: - - name: results - mount: /results +--- +apiVersion: batch/v1 +kind: Job +metadata: + name: benchmark-client + namespace: coldpress-project +spec: + template: + metadata: + labels: + app: benchmark-client + spec: + restartPolicy: Never + tolerations: + - operator: Exists + containers: + - name: client + image: ghcr.io/vllm-project/guidellm:nightly + command: + - guidellm + - benchmark + - run + - --target=http://localhost:8000 + - --output-dir=/results + - --outputs=json + - --max-seconds=30 + - --rate-type=throughput + - --rate=1 + - --data=prompt_tokens=256,output_tokens=128 + env: + - name: HOME + value: /tmp + volumeMounts: + - name: results + mountPath: /results + volumes: + - name: results + persistentVolumeClaim: + claimName: coldpress-project-storage diff --git a/examples/vllm_guidellm_benchmark/run.sh b/examples/vllm_guidellm_benchmark/run.sh new file mode 100755 index 0000000..d8d8d44 --- /dev/null +++ b/examples/vllm_guidellm_benchmark/run.sh @@ -0,0 +1,103 @@ +#!/bin/bash +# Manual Kubernetes orchestration - coordinate server and client + +set -e + +NAMESPACE="coldpress-project" + +echo "=========================================" +echo "vLLM + GuideLLM Benchmark Workflow" +echo "=========================================" +echo "" + +# Step 1: Apply server and service +echo "Step 1: Deploying vLLM inference server..." +oc apply -f vllm-server.yaml +oc apply -f service.yaml + +# Step 2: Wait for pod to be created +echo "Step 2: Waiting for server pod to be created..." +sleep 5 + +# Get pod name +POD_NAME=$(oc get pods -n $NAMESPACE -l app=vllm-server --no-headers -o custom-columns=":metadata.name" | head -1) + +if [ -z "$POD_NAME" ]; then + echo "ERROR: Server pod not found!" + exit 1 +fi + +echo "Server pod: $POD_NAME" + +# Step 3: Wait for pod to be ready (this is the hard part!) +echo "Step 3: Waiting for server to be ready..." +echo " - Checking readiness probe..." + +MAX_WAIT=600 # 10 minutes +ELAPSED=0 +READY=false + +while [ $ELAPSED -lt $MAX_WAIT ]; do + # Check if pod is ready + READY_STATUS=$(oc get pod $POD_NAME -n $NAMESPACE -o jsonpath='{.status.conditions[?(@.type=="Ready")].status}' 2>/dev/null || echo "Unknown") + + if [ "$READY_STATUS" == "True" ]; then + echo " ✓ Server is ready!" + READY=true + break + fi + + # Show progress + if [ $((ELAPSED % 10)) -eq 0 ]; then + PHASE=$(oc get pod $POD_NAME -n $NAMESPACE -o jsonpath='{.status.phase}' 2>/dev/null || echo "Unknown") + echo " - Still waiting... (${ELAPSED}s elapsed, phase: $PHASE)" + fi + + sleep 2 + ELAPSED=$((ELAPSED + 2)) +done + +if [ "$READY" != "true" ]; then + echo "ERROR: Server never became ready after ${MAX_WAIT}s" + echo "Pod status:" + oc describe pod $POD_NAME -n $NAMESPACE + exit 1 +fi + +# Step 4: Verify service endpoint is accessible +echo "Step 4: Verifying service endpoint..." +echo " - Testing http://vllm-server:8000/health" + +# We need a test pod to check from inside the cluster +TEST_POD="endpoint-test-$$" +oc run $TEST_POD -n $NAMESPACE --image=curlimages/curl:latest --rm -i --restart=Never --command -- \ + curl -s http://vllm-server:8000/health > /dev/null + +if [ $? -eq 0 ]; then + echo " ✓ Service endpoint is accessible!" +else + echo "ERROR: Service endpoint not accessible" + exit 1 +fi + +# Step 5: Launch benchmark client +echo "Step 5: Launching benchmark client..." +oc apply -f guidellm-client.yaml + +# Step 6: Wait for client to complete +echo "Step 6: Waiting for benchmark to complete..." +oc wait --for=condition=complete job/guidellm-benchmark -n $NAMESPACE --timeout=5m || { + echo "Benchmark job status:" + oc get job guidellm-benchmark -n $NAMESPACE + oc get pods -n $NAMESPACE -l app=guidellm-client + exit 1 +} + +echo "" +echo "=========================================" +echo "✓ Benchmark workflow completed!" +echo "=========================================" +echo "" +echo "To cleanup:" +echo " oc delete job vllm-inference-server guidellm-benchmark -n $NAMESPACE" +echo " oc delete service vllm-server -n $NAMESPACE" diff --git a/examples/vllm_guidellm_benchmark/service.yaml b/examples/vllm_guidellm_benchmark/service.yaml new file mode 100644 index 0000000..e7777a3 --- /dev/null +++ b/examples/vllm_guidellm_benchmark/service.yaml @@ -0,0 +1,13 @@ +apiVersion: v1 +kind: Service +metadata: + name: vllm-server + namespace: coldpress-project +spec: + selector: + app: vllm-server + ports: + - protocol: TCP + port: 8000 + targetPort: 8000 + name: http diff --git a/examples/vllm_guidellm_benchmark/vllm-server.yaml b/examples/vllm_guidellm_benchmark/vllm-server.yaml new file mode 100644 index 0000000..35c11ff --- /dev/null +++ b/examples/vllm_guidellm_benchmark/vllm-server.yaml @@ -0,0 +1,57 @@ +apiVersion: batch/v1 +kind: Job +metadata: + name: vllm-inference-server + namespace: coldpress-project +spec: + template: + metadata: + labels: + app: vllm-server + spec: + restartPolicy: Never + tolerations: + - operator: Exists + containers: + - name: server + image: nvcr.io/nvidia/vllm:26.03-py3 + command: ["bash", "-c"] + args: + - | + python -m vllm.entrypoints.openai.api_server \ + --model="${VLLM_MODEL}" \ + --port="${VLLM_PORT}" \ + --max-model-len="${VLLM_MAX_MODEL_LEN}" \ + --gpu-memory-utilization="${VLLM_GPU_MEMORY_UTIL}" + env: + - name: HOME + value: /tmp + - name: VLLM_MODEL + value: "ibm-granite/granite-3.3-8b-instruct" + - name: VLLM_PORT + value: "8000" + - name: VLLM_MAX_MODEL_LEN + value: "10000" + - name: VLLM_GPU_MEMORY_UTIL + value: "0.6" + ports: + - containerPort: 8000 + name: http + readinessProbe: + httpGet: + path: /health + port: 8000 + initialDelaySeconds: 30 + periodSeconds: 10 + resources: + requests: + nvidia.com/gpu: "1" + limits: + nvidia.com/gpu: "1" + volumeMounts: + - name: results + mountPath: /results + volumes: + - name: results + persistentVolumeClaim: + claimName: coldpress-project-storage diff --git a/projects/coldpress-project.yaml b/projects/coldpress-project.yaml index c2c2e08..76fede9 100644 --- a/projects/coldpress-project.yaml +++ b/projects/coldpress-project.yaml @@ -1,6 +1,10 @@ # Generated by: Claude Sonnet 4.5 namespace: coldpress-project +# Target platforms - determines which RBAC permissions to grant +# Options: jobset, kubeflow, kuberay, or combination (comma-separated) +targets: jobset,kubeflow,kuberay + # References to cluster-wide resources cluster_queue: coldpress-cluster-queue storage_class: nfs-csi diff --git a/pyproject.toml b/pyproject.toml index 82e0507..d0629ff 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,10 +5,10 @@ build-backend = "hatchling.build" [project] name = "coldpress" -version = "0.2.0" +version = "0.2.1" description = "AI/HPC workload orchestration for Kubernetes/OpenShift" readme = "README.md" -requires-python = ">=3.9" +requires-python = ">=3.10" license = {text = "See LICENSE file"} dependencies = [ "click>=8.0.0", diff --git a/tests/README.md b/tests/README.md deleted file mode 100644 index c4ad69b..0000000 --- a/tests/README.md +++ /dev/null @@ -1,119 +0,0 @@ -# Coldpress Test Suite - -This directory contains the comprehensive test suite for Coldpress. - -## Test Files - -### Validation Tests -**`test_validation.py`** - Tests Pydantic model validation -- Validates config.yaml schema -- Validates job-spec.yaml schema -- Validates project config schema -- Tests error handling for invalid configs -- Verifies validation catches errors early - -### Label Tests -**`test_labels.py`** - Tests standard Kubernetes labels -- Verifies all resources have `app.kubernetes.io/managed-by: coldpress` -- Verifies version labels -- Tests job-specific labels (`coldpress.io/job-id`) -- Validates labels on: JobSet, Service, ConfigMap, Namespace, LocalQueue, PVC, RBAC - -### Security Tests -**`test_security.py`** - Tests security improvements -- Tests JSON injection prevention (json.dumps vs f-strings) -- Tests Kubernetes name validation (lowercase, alphanumeric, dashes, dots) -- Tests Pydantic validators enforce naming rules -- Audits for hardcoded temp file vulnerabilities -- Validates user input sanitization - -### Error Handling Tests -**`test_error_handling.py`** - Tests error handling -- Tests exit code propagation (0/1/2) -- Tests specific exception types (FileNotFoundError, yaml.YAMLError, etc.) -- Tests validation error messages -- Verifies Click command error handling - -### Exit Code Tests -**`test_exit_codes.sh`** - Shell script tests for exit codes -- Tests success case (exit 0) -- Tests error cases (exit 1) -- Tests usage errors (exit 2) -- Verifies shell-level exit code propagation - -### RoCE Disabled Tests -**`test_roce_disabled.py`** - Tests RoCE NIC support is disabled -- Verifies no RDMA resources in ClusterQueue -- Verifies no NetworkAttachmentDefinitions generated -- Verifies `roce_nics` field still validates in config -- Tests simplified cluster setup - -## Running Tests - -### Run All Tests -```bash -# From repository root -./tests/run_all_tests.sh -``` - -### Run Individual Tests -```bash -# Validation -python tests/test_validation.py - -# Labels -python tests/test_labels.py - -# Security -python tests/test_security.py - -# Error handling -python tests/test_error_handling.py - -# RoCE disabled -python tests/test_roce_disabled.py - -# Exit codes -bash tests/test_exit_codes.sh -``` - -### GitHub Actions -Tests run automatically on every push and pull request via GitHub Actions (`.github/workflows/tests.yml`). - -The workflow runs all tests on Python 3.9, 3.10, 3.11, 3.12, 3.13, and 3.14. - -## Test Coverage - -These tests verify fixes for GitHub issue #37: -- ✅ Issue #1: No Input Validation → `test_validation.py` -- ✅ Issue #2: Poor Separation of Concerns → Validated by integration tests -- ✅ Issue #3: Error Handling → `test_error_handling.py`, `test_exit_codes.sh` -- ✅ Issue #6: Security → `test_security.py` -- ✅ Issue #7: Error Handling → `test_error_handling.py` -- ✅ Issue #8: No Consistent Labeling → `test_labels.py` - -## Adding New Tests - -When adding new tests: -1. Create `test_*.py` in this directory -2. Add to `run_all_tests.sh` -3. Add to `.github/workflows/tests.yml` -4. Update this README - -## Dependencies - -Tests require: -- Python 3.9+ -- pyyaml -- click -- pydantic - -Install with: -```bash -pip install pyyaml click pydantic -``` - -Or use the development environment: -```bash -source .venv/bin/activate -``` diff --git a/tests/test_error_handling.py b/tests/test_error_handling.py index facc7d6..a52d93e 100755 --- a/tests/test_error_handling.py +++ b/tests/test_error_handling.py @@ -19,8 +19,8 @@ def test_exit_codes(): tests = [ { - "name": "Success: Valid coldpress config", - "cmd": "python -m coldpress.cli generate --config examples/pytorch_ddp_training/config.yaml", + "name": "Success: Valid coldpress intent", + "cmd": "python -m coldpress.cli generate --intent examples/pytorch_ddp_training/intent_jobset.yaml", "expected": 0, }, { @@ -68,13 +68,7 @@ def test_specific_exceptions(): print("=" * 60) # Check that code uses specific exceptions - import coldpress.generator as gen - - # Test create_service with invalid URL - task = {"health_check": None} - result = gen.create_service(task, 0, "test", "coldpress-test", "default") - assert result is None, "create_service should return None for missing health_check" - print("✅ create_service handles missing health_check") + import coldpress.jobset_generator as gen # Test build_discovery_init_container with missing template result = gen.build_discovery_init_container( @@ -85,12 +79,14 @@ def test_specific_exceptions(): ) print("✅ build_discovery_init_container handles missing template") - # Test build_discovery_job with missing template - result = gen.build_discovery_job( - "/tmp/non-existent-template.yaml", "base", "pvc", "0" - ) - assert result is None, "build_discovery_job should return None for missing template" - print("✅ build_discovery_job handles missing template") + # Test substitute_macros with dangerous characters + try: + macros = {"TEST": "value; rm -rf /"} + gen.substitute_macros("${TEST}", macros) + raise AssertionError("Should have caught dangerous characters in macro value") + except ValueError as e: + print("✅ substitute_macros rejects dangerous characters") + assert "Invalid character" in str(e) def test_validation_errors(): diff --git a/tests/test_exit_codes.sh b/tests/test_exit_codes.sh index 2b82a4f..b37f24e 100755 --- a/tests/test_exit_codes.sh +++ b/tests/test_exit_codes.sh @@ -17,10 +17,10 @@ else echo "❌ Exit code $EXIT_CODE (expected 1)" fi -# Test 2: coldpress with missing config (should return 2 for usage error) +# Test 2: coldpress with missing intent file (should return 2 for usage error) echo "" -echo "Test 2: coldpress with non-existent config..." -python -m coldpress.cli generate --config non-existent.yaml >/dev/null 2>&1 +echo "Test 2: coldpress with non-existent intent..." +python -m coldpress.cli generate --intent non-existent.yaml >/dev/null 2>&1 EXIT_CODE=$? if [ $EXIT_CODE -eq 2 ]; then echo "✅ Exit code 2 (expected for Click usage error)" @@ -28,10 +28,10 @@ else echo "❌ Exit code $EXIT_CODE (expected 2)" fi -# Test 3: coldpress with valid config (should return 0) +# Test 3: coldpress with valid intent (should return 0) echo "" -echo "Test 3: coldpress with valid config..." -python -m coldpress.cli generate --config examples/pytorch_ddp_training/config.yaml >/dev/null 2>&1 +echo "Test 3: coldpress with valid intent..." +python -m coldpress.cli generate --intent examples/pytorch_ddp_training/intent_jobset.yaml >/dev/null 2>&1 EXIT_CODE=$? if [ $EXIT_CODE -eq 0 ]; then echo "✅ Exit code 0 (expected for success)" diff --git a/tests/test_labels.py b/tests/test_labels.py index 08b295c..a897f09 100644 --- a/tests/test_labels.py +++ b/tests/test_labels.py @@ -1,8 +1,10 @@ #!/usr/bin/env python3 """Test script to verify Coldpress resource labels are applied correctly.""" -from coldpress.generator import ( - generate_jobset, +from coldpress.jobset_generator import ( + generate_jobset_from_intent, +) +from coldpress.constants import ( COLDPRESS_LABELS as COLDPRESS_JOB_LABELS, ) from coldpress_setup.generator import ( @@ -12,6 +14,7 @@ generate_cluster_queue, COLDPRESS_LABELS as COLDPRESS_SETUP_LABELS, ) +from coldpress_common import validate_intent def test_label_constants(): @@ -22,7 +25,7 @@ def test_label_constants(): expected = { "app.kubernetes.io/managed-by": "coldpress", - "app.kubernetes.io/version": "0.2.0", + "app.kubernetes.io/version": "0.2.1", } assert COLDPRESS_JOB_LABELS == expected, ( @@ -40,70 +43,94 @@ def test_jobset_labels(): print("Testing JobSet Labels") print("=" * 60) - job_spec = { - "name": "test-job", - "namespace": "test-ns", - "tasks": [ - { - "name": "task-0", - "containers": [{"name": "main", "image": "alpine:latest"}], + # Create vanilla k8s Job + vk8s_job = { + "apiVersion": "batch/v1", + "kind": "Job", + "metadata": {"name": "test-task"}, + "spec": { + "template": { + "spec": { + "containers": [{"name": "main", "image": "alpine:latest"}], + "restartPolicy": "Never", + } } - ], - "storage": {"results": "test-pvc"}, + }, + } + + # Create intent config + intent_data = { + "project": "coldpress-project", + "output": "test-job", + "target": "jobset", + "tasks": [{"name": "test-task", "replicas": 1}], } + intent_config = validate_intent(intent_data) + + # Create project config + project_config = {"namespace": "test-ns", "storage": {"results": "test-pvc"}} - jobset, services, _ = generate_jobset(job_spec, {0: "any"}) + jobset, services, _ = generate_jobset_from_intent( + {"test-task": vk8s_job}, intent_config, project_config, "test-ns" + ) labels = jobset["metadata"]["labels"] assert "app.kubernetes.io/managed-by" in labels assert labels["app.kubernetes.io/managed-by"] == "coldpress" assert "app.kubernetes.io/version" in labels - assert labels["app.kubernetes.io/version"] == "0.2.0" - assert "coldpress.io/job-id" in labels - assert labels["coldpress.io/job-id"] == "coldpress-test-job" + assert labels["app.kubernetes.io/version"] == "0.2.1" print(f"✅ JobSet labels: {labels}") def test_service_labels(): - """Test that Services have correct labels.""" + """Test that Services are not generated (JobSet provides automatic DNS).""" print("\n" + "=" * 60) - print("Testing Service Labels") + print("Testing Service Labels (JobSet DNS)") print("=" * 60) - job_spec = { - "name": "test-job", - "namespace": "test-ns", - "tasks": [ - { - "name": "server", - "blocking": "endpoint", - "health_check": "http://server:8000/health", - "containers": [ - { - "name": "main", - "image": "vllm:latest", - "readinessProbe": { - "httpGet": {"path": "/health", "port": 8000} - }, - } - ], + # Create vanilla k8s Job with readinessProbe + vk8s_job = { + "apiVersion": "batch/v1", + "kind": "Job", + "metadata": {"name": "server"}, + "spec": { + "template": { + "spec": { + "containers": [ + { + "name": "main", + "image": "vllm:latest", + "readinessProbe": { + "httpGet": {"path": "/health", "port": 8000} + }, + } + ], + "restartPolicy": "Never", + } } - ], - "storage": {"results": "test-pvc"}, + }, } - jobset, services, _ = generate_jobset(job_spec, {0: "any"}) + # Create intent config + intent_data = { + "project": "coldpress-project", + "output": "test-job", + "target": "jobset", + "tasks": [{"name": "server", "replicas": 1}], + } + intent_config = validate_intent(intent_data) - assert len(services) > 0, "Expected at least one service" - labels = services[0]["metadata"]["labels"] + # Create project config + project_config = {"namespace": "test-ns", "storage": {"results": "test-pvc"}} - assert "app.kubernetes.io/managed-by" in labels - assert labels["app.kubernetes.io/managed-by"] == "coldpress" - assert "app.kubernetes.io/version" in labels - assert "coldpress.io/job-id" in labels + jobset, services, _ = generate_jobset_from_intent( + {"server": vk8s_job}, intent_config, project_config, "test-ns" + ) - print(f"✅ Service labels: {labels}") + # Services are no longer generated - JobSet provides automatic DNS + assert len(services) == 0, "Services should not be generated (JobSet provides DNS)" + print("✅ No services generated (JobSet provides automatic DNS)") def test_project_resource_labels(): @@ -206,7 +233,7 @@ def main(): print("=" * 60) print("\nAll Coldpress resources have standard labels:") print(" - app.kubernetes.io/managed-by: coldpress") - print(" - app.kubernetes.io/version: 0.2.0") + print(" - app.kubernetes.io/version: 0.2.1") print(" - coldpress.io/job-id: {job_name} (for job resources)") print("\nQuery all resources:") print( diff --git a/tests/test_script_gen_security.py b/tests/test_script_gen_security.py new file mode 100644 index 0000000..3a19937 --- /dev/null +++ b/tests/test_script_gen_security.py @@ -0,0 +1,245 @@ +#!/usr/bin/env python3 +"""Test script generation security - verify filenames are sanitized.""" + +import sys +import pytest +from coldpress.script_gen import ( + generate_run_script, + sanitize_filename, + sanitize_identifier, +) + + +def test_sanitize_identifier(): + """Test identifier (job name, namespace) sanitization.""" + print("\n" + "=" * 60) + print("Testing Identifier Sanitization (job names, namespaces)") + print("=" * 60) + + # Valid identifiers should pass + valid_cases = [ + "test-job", + "my-namespace", + "job-123", + "coldpress-test", + "prod-cluster", + ] + + for identifier in valid_cases: + try: + result = sanitize_identifier(identifier, "job name") + assert result == identifier + print(f"✅ Valid: '{identifier}' -> '{result}'") + except ValueError as e: + pytest.fail(f"Valid identifier rejected: {identifier} - {e}") + + # Invalid identifiers should fail + invalid_cases = [ + ("job;rm -rf /", "command injection"), + ("$(whoami)", "command substitution"), + ("`id`", "backtick substitution"), + ("job|cat", "pipe"), + ("job&background", "background"), + ("job>output", "redirection"), + ("job name", "space"), + ("'job'", "single quote"), + ('"job"', "double quote"), + ("/etc/passwd", "path separator"), + ("..\\windows", "backslash"), + ("job$VAR", "dollar sign"), + ] + + for identifier, reason in invalid_cases: + try: + result = sanitize_identifier(identifier, "job name") + pytest.fail( + f"Invalid identifier should have been rejected: {identifier} ({reason})" + ) + except ValueError: + print(f"✅ Rejected: '{identifier}' ({reason})") + + +def test_sanitize_filename_edge_cases(): + """Test edge cases for filename sanitization.""" + print("\n" + "=" * 60) + print("Testing Filename Sanitization Edge Cases") + print("=" * 60) + + # Valid filenames should pass + valid_cases = [ + "config.yaml", + "script_v2.sh", + "data.tar.gz", + "README.md", + "file-123.txt", + ] + + for filename in valid_cases: + try: + result = sanitize_filename(filename) + assert result == filename + print(f"✅ Valid: '{filename}' -> '{result}'") + except ValueError as e: + pytest.fail(f"Valid filename rejected: {filename} - {e}") + + # Invalid filenames should fail + invalid_cases = [ + ("../../../etc/passwd", "path traversal"), + ("file;rm -rf /", "command injection"), + ("$(whoami).txt", "command substitution"), + ("`id`.txt", "backtick substitution"), + ("file|cat", "pipe"), + ("file&background", "background"), + ("file>output", "redirection"), + ("file name.txt", "space"), + ("'file'.txt", "single quote"), + ('"file".txt', "double quote"), + ] + + for filename, reason in invalid_cases: + try: + result = sanitize_filename(filename) + pytest.fail( + f"Invalid filename should have been rejected: {filename} ({reason})" + ) + except ValueError: + print(f"✅ Rejected: '{filename}' ({reason})") + + +def test_generate_run_script_with_safe_files(): + """Test that run script generation works with safe filenames.""" + print("\n" + "=" * 60) + print("Testing Run Script Generation with Safe Files") + print("=" * 60) + + safe_files = ["config.yaml", "script.sh", "data.json"] + + try: + script = generate_run_script( + job_name="test-job", + namespace="test-ns", + configmap_name="test-config", + configmap_files=safe_files, + manifest_type="jobset", + ) + + # Verify the script contains the expected files + assert "--from-file=config.yaml" in script + assert "--from-file=script.sh" in script + assert "--from-file=data.json" in script + assert "coldpress-test-config" in script + print("✅ Run script generated successfully with safe files") + print(f" Files: {', '.join(safe_files)}") + except Exception as e: + pytest.fail(f"Failed to generate script with safe files: {e}") + + +def test_generate_run_script_with_dangerous_files(): + """Test that run script generation rejects dangerous filenames.""" + print("\n" + "=" * 60) + print("Testing Run Script Generation with Dangerous Files") + print("=" * 60) + + dangerous_test_cases = [ + (["config.yaml", "file;rm -rf /"], "command injection"), + (["../../../etc/passwd"], "path traversal"), + (["$(whoami).txt"], "command substitution"), + (["`id`.txt"], "backtick substitution"), + (["file|cat"], "pipe character"), + (["file name.txt"], "space in filename"), + ] + + for files, reason in dangerous_test_cases: + try: + generate_run_script( + job_name="test-job", + namespace="test-ns", + configmap_name="test-config", + configmap_files=files, + manifest_type="jobset", + ) + pytest.fail( + f"Script generation should have rejected dangerous files: {files} ({reason})" + ) + except ValueError as e: + print(f"✅ Rejected: {files} ({reason})") + assert "Invalid" in str(e) or "security" in str(e).lower() + + +def test_no_injection_in_generated_script(): + """Test that user-provided filenames don't inject into generated scripts.""" + print("\n" + "=" * 60) + print("Testing Generated Scripts for User Input Injection") + print("=" * 60) + + # Generate a script with safe files + safe_files = ["config.yaml", "app.json"] + script = generate_run_script( + job_name="test-job", + namespace="test-ns", + configmap_name="test-config", + configmap_files=safe_files, + manifest_type="jobset", + ) + + # Check that user filenames appear safely in the script + # They should appear as --from-file= with no dangerous characters + assert "--from-file=config.yaml" in script + assert "--from-file=app.json" in script + + # Check that the ConfigMap section doesn't contain dangerous user input patterns + # Extract just the ConfigMap creation line + configmap_section = "" + for line in script.split("\n"): + if "oc create configmap" in line: + configmap_section = line + break + + # Verify filenames in the configmap command are safe + assert "../" not in configmap_section # No path traversal + assert ";rm" not in configmap_section # No command injection + assert "|cat" not in configmap_section # No piping + assert "&background" not in configmap_section # No background execution + assert ">output" not in configmap_section # No redirection + # Note: $(dirname "$0") is legitimate bash and appears in the template, not user input + + print("✅ User-provided filenames safely embedded in generated script") + print(f" Safe files: {', '.join(safe_files)}") + + +def main(): + """Run all script generation security tests.""" + print("\n" + "=" * 60) + print("SCRIPT GENERATION SECURITY TEST SUITE") + print("=" * 60) + + try: + test_sanitize_identifier() + test_sanitize_filename_edge_cases() + test_generate_run_script_with_safe_files() + test_generate_run_script_with_dangerous_files() + test_no_injection_in_generated_script() + + print("\n" + "=" * 60) + print("✅ All script generation security tests passed!") + print("=" * 60) + print("\nSecurity improvements:") + print(" ✅ Job names and namespaces sanitized") + print(" ✅ Filenames sanitized before use in shell commands") + print(" ✅ Path traversal attempts blocked") + print(" ✅ Command injection attempts blocked") + print(" ✅ Shell metacharacters rejected") + print("=" * 60) + return 0 + except (AssertionError, Exception) as e: + print("\n" + "=" * 60) + print(f"❌ Script generation security test failed: {e}") + print("=" * 60) + import traceback + + traceback.print_exc() + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_security.py b/tests/test_security.py index 29038f3..936517c 100755 --- a/tests/test_security.py +++ b/tests/test_security.py @@ -1,10 +1,8 @@ #!/usr/bin/env python3 """Test security and reliability improvements in Coldpress.""" -import json import sys import pytest -from coldpress_setup.generator import generate_sriov_network_attachments from coldpress_common import ( validate_kubernetes_name, validate_project_config, @@ -14,41 +12,6 @@ from pydantic import ValidationError -def test_json_injection_fix(): - """Test that JSON is constructed safely using json.dumps.""" - print("=" * 60) - print("Testing JSON Injection Prevention") - print("=" * 60) - - # Generate SRIOV network attachments - attachments = generate_sriov_network_attachments("test-namespace", 2) - - passed = 0 - failed = 0 - - for attachment in attachments: - config_str = attachment["spec"]["config"] - - # Check that config is valid JSON - try: - config = json.loads(config_str) - print(f"✅ Valid JSON: {config['name']}") - passed += 1 - except json.JSONDecodeError: - print(f"❌ Invalid JSON in config: {config_str}") - failed += 1 - - # Verify structure - if "cniVersion" in config and "type" in config and "ipam" in config: - print(" ✅ Correct CNI config structure") - else: - print(" ❌ Missing required CNI config fields") - failed += 1 - - print(f"\nJSON tests: {passed} passed, {failed} failed") - assert failed == 0, f"JSON injection tests failed: {failed} failures" - - def test_kubernetes_name_validation(): """Test that Kubernetes names are validated.""" print("\n" + "=" * 60) @@ -233,6 +196,59 @@ def test_no_temp_file_vulnerabilities(): print(" (Container mount paths like /tmp/result are safe)") +def test_filename_sanitization(): + """Test that filenames are sanitized before use in shell commands.""" + print("\n" + "=" * 60) + print("Testing Filename Sanitization") + print("=" * 60) + + from coldpress.script_gen import sanitize_filename + + test_cases = [ + # (filename, should_pass, description) + ("valid-file.txt", True, "Valid filename"), + ("file_123.yaml", True, "Valid with underscore and number"), + ("file.tar.gz", True, "Valid with multiple dots"), + ("../etc/passwd", False, "Path traversal attempt"), + ("dir/file.txt", False, "Directory separator"), + ("file;rm -rf /", False, "Command injection with semicolon"), + ("file`whoami`.txt", False, "Command substitution with backticks"), + ("file$(whoami).txt", False, "Command substitution with $()"), + ("file|cat", False, "Pipe character"), + ("file&background", False, "Background execution"), + ("file>output", False, "Redirection"), + ("file name.txt", False, "Space in filename"), + ("file'quote.txt", False, "Single quote"), + ('file"quote.txt', False, "Double quote"), + ("", False, "Empty string"), + (".", False, "Current directory"), + ("..", False, "Parent directory"), + ] + + passed = 0 + failed = 0 + + for filename, should_pass, description in test_cases: + try: + sanitized = sanitize_filename(filename) + if should_pass: + print(f"✅ {description}: '{filename}' accepted as '{sanitized}'") + passed += 1 + else: + print(f"❌ {description}: '{filename}' should have been rejected") + failed += 1 + except ValueError as e: + if not should_pass: + print(f"✅ {description}: '{filename}' rejected") + passed += 1 + else: + print(f"❌ {description}: '{filename}' should have been accepted - {e}") + failed += 1 + + print(f"\nFilename sanitization tests: {passed} passed, {failed} failed") + assert failed == 0, f"Filename sanitization tests failed: {failed} failures" + + def main(): """Run all security tests.""" print("\n" + "=" * 60) @@ -240,8 +256,8 @@ def main(): print("=" * 60) try: - # Test JSON injection prevention - test_json_injection_fix() + # YAML is constructed with yaml.safe_dump() (verified in code) + # No separate test needed - safe serialization is used throughout # Test Kubernetes name validation test_kubernetes_name_validation() @@ -252,14 +268,18 @@ def main(): # Test temp file security test_no_temp_file_vulnerabilities() + # Test filename sanitization + test_filename_sanitization() + print("\n" + "=" * 60) print("✅ All security tests passed!") print("=" * 60) print("\nSecurity improvements:") - print(" ✅ JSON constructed with json.dumps() (no injection)") + print(" ✅ YAML constructed with yaml.safe_dump() (no injection)") print(" ✅ Kubernetes names validated against spec") print(" ✅ User input sanitized before use in resource names") print(" ✅ No hardcoded temp file vulnerabilities") + print(" ✅ Filenames sanitized before use in shell commands") print("=" * 60) return 0 except (AssertionError, Exception) as e: diff --git a/tests/test_validation.py b/tests/test_validation.py index af13df0..1a573f9 100755 --- a/tests/test_validation.py +++ b/tests/test_validation.py @@ -3,8 +3,6 @@ import yaml from coldpress_common import ( - validate_config, - validate_task_specs, validate_project_config, ) from pydantic import ValidationError @@ -16,23 +14,24 @@ def test_valid_configs(): print("Testing VALID configurations") print("=" * 60) - # Test 1: Valid config.yaml - print("\n1. Testing config.yaml validation...") - with open("examples/pytorch_ddp_training/config.yaml") as f: - config_data = yaml.safe_load(f) - config = validate_config(config_data) - print( - f" ✓ Valid config: project={config.project}, discovery={config.discovery}" - ) - - # Test 2: Valid job-spec.yaml - print("\n2. Testing job-spec.yaml validation...") + # Test 1: Valid intent.yaml + print("\n1. Testing intent.yaml validation...") + from coldpress_common import validate_intent + + with open("examples/pytorch_ddp_training/intent_jobset.yaml") as f: + intent_data = yaml.safe_load(f) + intent = validate_intent(intent_data) + print(f" ✓ Valid intent: target={intent.target}, tasks={len(intent.tasks)}") + + # Test 2: Valid job-spec.yaml (vanilla k8s Jobs) + print("\n2. Testing job-spec.yaml (vanilla k8s)...") with open("examples/pytorch_ddp_training/job-spec.yaml") as f: - task_specs = list(yaml.safe_load_all(f)) - validated = validate_task_specs(task_specs) - print(f" ✓ Valid task spec: {len(validated)} task(s)") - for i, task in enumerate(validated): - print(f" - Task {i}: {task.name} ({len(task.containers)} container(s))") + manifests = list(yaml.safe_load_all(f)) + jobs = [m for m in manifests if m.get("kind") == "Job"] + print(f" ✓ Valid job spec: {len(jobs)} Job(s)") + for i, job in enumerate(jobs): + job_name = job["metadata"]["name"] + print(f" - Job {i}: {job_name}") # Test 3: Valid project config print("\n3. Testing project config validation...") @@ -41,15 +40,14 @@ def test_valid_configs(): project = validate_project_config(project_data) print(f" ✓ Valid project: namespace={project.namespace}") - # Test 4: Multi-task job spec - print("\n4. Testing multi-task job-spec.yaml...") - with open("examples/vllm_guidellm_benchmark/job-spec.yaml") as f: - task_specs = list(yaml.safe_load_all(f)) - validated = validate_task_specs(task_specs) - print(f" ✓ Valid multi-task spec: {len(validated)} task(s)") - for i, task in enumerate(validated): - blocking = task.blocking or "completion" - print(f" - Task {i}: {task.name} (blocking={blocking})") + # Test 4: Multi-task intent spec + print("\n4. Testing multi-task intent.yaml...") + with open("examples/vllm_guidellm_benchmark/intent_jobset.yaml") as f: + intent_data = yaml.safe_load(f) + intent = validate_intent(intent_data) + print(f" ✓ Valid multi-task intent: {len(intent.tasks)} task(s)") + for i, task in enumerate(intent.tasks): + print(f" - Task {i}: {task.name} (replicas={task.replicas})") def test_invalid_configs(): @@ -58,51 +56,53 @@ def test_invalid_configs(): print("Testing INVALID configurations (should catch errors)") print("=" * 60) - # Test 1: Missing required field - print("\n1. Testing missing 'name' field...") + from coldpress_common import validate_intent + + # Test 1: Missing required field (project) + print("\n1. Testing missing 'project' field...") try: - invalid_task = [{"containers": [{"name": "test", "image": "alpine"}]}] - validate_task_specs(invalid_task) - print(" ❌ ERROR: Should have caught missing 'name' field!") + invalid_intent = { + "target": "jobset", + "output": "test", + "tasks": [{"name": "test-task", "replicas": 1}], + } + validate_intent(invalid_intent) + print(" ❌ ERROR: Should have caught missing 'project' field!") except (ValidationError, ValueError) as e: print(f" ✓ Caught error: {str(e).splitlines()[0][:70]}...") - # Test 2: Missing containers - print("\n2. Testing missing 'containers' field...") + # Test 2: Missing tasks + print("\n2. Testing missing 'tasks' field...") try: - invalid_task = [{"name": "test-task"}] - validate_task_specs(invalid_task) - print(" ❌ ERROR: Should have caught missing 'containers' field!") + invalid_intent = {"project": "test", "output": "test", "target": "jobset"} + validate_intent(invalid_intent) + print(" ❌ ERROR: Should have caught missing 'tasks' field!") except (ValidationError, ValueError) as e: print(f" ✓ Caught error: {str(e).splitlines()[0][:70]}...") - # Test 3: Invalid blocking endpoint without health check - print("\n3. Testing endpoint blocking without health check...") + # Test 3: Invalid target type + print("\n3. Testing invalid target type...") try: - invalid_task = [ - { - "name": "test-task", - "blocking": "endpoint", - "containers": [{"name": "test", "image": "alpine"}], - } - ] - validate_task_specs(invalid_task) - print(" ❌ ERROR: Should have caught missing health check!") + invalid_intent = { + "project": "test", + "output": "test", + "target": "invalid_target", + "tasks": [{"name": "test", "replicas": 1}], + } + validate_intent(invalid_intent) + print(" ❌ ERROR: Should have caught invalid target type!") except (ValidationError, ValueError) as e: print(f" ✓ Caught error: {str(e).splitlines()[0][:70]}...") - # Test 4: Invalid blocking type - print("\n4. Testing invalid blocking type...") + # Test 4: Invalid replicas (negative) + print("\n4. Testing invalid replicas...") try: - invalid_task = [ - { - "name": "test-task", - "blocking": "invalid_type", - "containers": [{"name": "test", "image": "alpine"}], - } - ] - validate_task_specs(invalid_task) - print(" ❌ ERROR: Should have caught invalid blocking type!") + invalid_intent = { + "backend": "jobset", + "tasks": [{"name": "test", "replicas": -1}], + } + validate_intent(invalid_intent) + print(" ❌ ERROR: Should have caught invalid replicas!") except (ValidationError, ValueError) as e: print(f" ✓ Caught error: {str(e).splitlines()[0][:70]}...")