From 251eb384928b4513bd15e49a1b69cc0c8e8c6f22 Mon Sep 17 00:00:00 2001 From: asanaullah Date: Thu, 16 Apr 2026 08:42:42 -0700 Subject: [PATCH 01/14] added kubeflow and kserve support; added intent based gen --- .gitignore | 12 + README.md | 414 ------- cluster/ocp-test-nerc-mghpcc.yaml | 8 +- coldpress/cli.py | 406 +++++-- coldpress/constants.py | 39 + coldpress/generator.py | 851 ++----------- coldpress/jobset_generator.py | 678 +++++++++++ coldpress/jobset_generator_legacy.py | 803 ++++++++++++ coldpress/kubeflow_generator.py | 1076 +++++++++++++++++ coldpress/script_gen.py | 263 +++- coldpress_common/__init__.py | 8 + coldpress_common/model.py | 118 +- coldpress_setup/cli.py | 23 +- coldpress_setup/generator.py | 85 +- discovery/README.md | 135 --- docs/CHANGELOG.md | 151 --- docs/ERROR_HANDLING.md | 202 ---- docs/LABELS.md | 100 -- docs/ROCE_DISABLED.md | 127 -- docs/TESTING.md | 228 ---- docs/VALIDATION.md | 215 ---- docs/quickstart_admin.md | 213 ---- docs/quickstart_user.md | 421 ------- examples/README.md | 180 --- .../pytorch_ddp_training/comp/job-spec.yaml | 75 ++ .../pytorch_ddp_training/comp/jobset.yaml | 596 +++++++++ .../pytorch_ddp_training/comp/pytorchjob.yaml | 131 ++ examples/pytorch_ddp_training/config.yaml | 9 - .../pytorch_ddp_training/intent_jobset.yaml | 21 + .../pytorch_ddp_training/intent_kubeflow.yaml | 21 + examples/pytorch_ddp_training/job-spec.yaml | 110 +- .../pytorch_ddp_training/model_config.json | 8 - examples/pytorch_ddp_training/train.py | 20 +- examples/pytorch_ddp_training_node1/README.md | 135 --- .../pytorch_ddp_training_node1/config.yaml | 18 - .../pytorch_ddp_training_node1/job-spec.yaml | 41 - .../model_config.json | 8 - examples/pytorch_ddp_training_node1/train.py | 170 --- .../guidellm-client.yaml | 52 + .../intent_jobset.yaml | 19 + .../intent_kserve.yaml | 14 + .../vllm_guidellm_benchmark/job-spec.yaml | 170 ++- examples/vllm_guidellm_benchmark/run.sh | 103 ++ examples/vllm_guidellm_benchmark/service.yaml | 13 + .../vllm_guidellm_benchmark/vllm-server.yaml | 57 + projects/coldpress-project.yaml | 4 + tests/README.md | 119 -- tests/test_error_handling.py | 26 +- tests/test_labels.py | 113 +- tests/test_validation.py | 118 +- 50 files changed, 4843 insertions(+), 4084 deletions(-) delete mode 100644 README.md create mode 100644 coldpress/constants.py create mode 100644 coldpress/jobset_generator.py create mode 100644 coldpress/jobset_generator_legacy.py create mode 100644 coldpress/kubeflow_generator.py delete mode 100644 discovery/README.md delete mode 100644 docs/CHANGELOG.md delete mode 100644 docs/ERROR_HANDLING.md delete mode 100644 docs/LABELS.md delete mode 100644 docs/ROCE_DISABLED.md delete mode 100644 docs/TESTING.md delete mode 100644 docs/VALIDATION.md delete mode 100644 docs/quickstart_admin.md delete mode 100644 docs/quickstart_user.md delete mode 100644 examples/README.md create mode 100644 examples/pytorch_ddp_training/comp/job-spec.yaml create mode 100644 examples/pytorch_ddp_training/comp/jobset.yaml create mode 100644 examples/pytorch_ddp_training/comp/pytorchjob.yaml delete mode 100644 examples/pytorch_ddp_training/config.yaml create mode 100644 examples/pytorch_ddp_training/intent_jobset.yaml create mode 100644 examples/pytorch_ddp_training/intent_kubeflow.yaml delete mode 100644 examples/pytorch_ddp_training/model_config.json delete mode 100644 examples/pytorch_ddp_training_node1/README.md delete mode 100644 examples/pytorch_ddp_training_node1/config.yaml delete mode 100644 examples/pytorch_ddp_training_node1/job-spec.yaml delete mode 100644 examples/pytorch_ddp_training_node1/model_config.json delete mode 100644 examples/pytorch_ddp_training_node1/train.py create mode 100644 examples/vllm_guidellm_benchmark/guidellm-client.yaml create mode 100644 examples/vllm_guidellm_benchmark/intent_jobset.yaml create mode 100644 examples/vllm_guidellm_benchmark/intent_kserve.yaml create mode 100755 examples/vllm_guidellm_benchmark/run.sh create mode 100644 examples/vllm_guidellm_benchmark/service.yaml create mode 100644 examples/vllm_guidellm_benchmark/vllm-server.yaml delete mode 100644 tests/README.md diff --git a/.gitignore b/.gitignore index fb0bcf0..6e3fbfe 100644 --- a/.gitignore +++ b/.gitignore @@ -52,6 +52,13 @@ dmypy.json # pyright .pyright/ +# Testing +.pytest_cache/ +.coverage +.coverage.* +htmlcov/ +.tox/ + # Jupyter Notebook .ipynb_checkpoints @@ -70,3 +77,8 @@ manifests/ # User-specific files users/test.yaml + +# Documentation +docs/ +README.md +CHANGELOG.md diff --git a/README.md b/README.md deleted file mode 100644 index 3fdf651..0000000 --- a/README.md +++ /dev/null @@ -1,414 +0,0 @@ - -# Coldpress - -[![Tests](https://github.com/asanaullah/coldpress/workflows/Tests/badge.svg?branch=v0.2)](https://github.com/asanaullah/coldpress/actions/workflows/tests.yml) - -Coldpress is a prescriptive manifest generator that reduces the effort and expertise needed to deploy complex AI/HPC workloads on Kubernetes clusters. - -**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? - -### Overall Flow - -``` -┌─────────────────────────────────────────────────────────────────┐ -│ Prerequisites: Kueue and JobSet operators must be installed │ -└─────────────────────────────────────────────────────────────────┘ - ↓ -┌─────────────────────────────────────────────────────────────────┐ -│ Phase 1: Admin Setup (One-time) │ -├─────────────────────────────────────────────────────────────────┤ -│ 1. coldpress-setup generate cluster → manifests/cluster-*.yaml │ -│ → manifests/label-nodes-*.sh │ -│ │ -│ 2. ./manifests/label-nodes-*.sh (labels nodes for scheduling) │ -│ │ -│ 3. oc apply -f manifests/cluster-*.yaml │ -│ │ -│ 4. coldpress-setup generate project → manifests/project-*.yaml │ -│ │ -│ 5. oc apply -f manifests/project-*.yaml │ -│ │ -│ 6. coldpress-setup generate user → manifests/user-*.yaml │ -│ │ -│ 7. oc apply -f manifests/user-*.yaml │ -│ │ -│ Creates: Node labels, ClusterQueue, ResourceFlavors, │ -│ Namespaces, PVCs, RBAC │ -└─────────────────────────────────────────────────────────────────┘ - ↓ -┌─────────────────────────────────────────────────────────────────┐ -│ Phase 2: User Workflow (Repeatable) │ -├─────────────────────────────────────────────────────────────────┤ -│ 1. coldpress generate --config job.yaml → output/job-name/ │ -│ - jobset.yaml (Kubernetes manifest) │ -│ - run.sh, monitor.sh, logs.sh, explore.sh, cp.sh, cleanup.sh │ -│ │ -│ 2. User reviews jobset.yaml │ -│ │ -│ 3. ./run.sh applies JobSet to cluster │ -│ │ -│ 4. Kueue schedules job when resources available │ -│ │ -│ 5. 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. ./explore.sh opens interactive shell to browse results │ -│ │ -│ 8. ./cp.sh copies results from PVC (optional) │ -│ │ -│ 9. ./cleanup.sh deletes JobSet (preserves results in PVC) │ -└─────────────────────────────────────────────────────────────────┘ -``` - -### What Coldpress Does - -**For Administrators (`coldpress-setup`):** -- Generates cluster-wide Kueue configuration (ClusterQueue, ResourceFlavors, LocalQueues) -- Generates project namespaces with PersistentVolumeClaims -- Generates user RBAC (RoleBindings) for job submission -- Generates node labeling scripts -- Outputs timestamped manifests for GitOps workflows - -**For Users (`coldpress`):** -- Generates JobSet manifests from simple job specifications -- 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 -- 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 - -Choose the installation method that fits your use case: - -| Use Case | Command | Activation Needed? | Best For | -|----------|---------|-------------------|----------| -| **Running jobs** | `./setup-env.sh --pipx` | ❌ No | End users, cluster users | -| **Running jobs + dev** | `./setup-env.sh --pipx-editable` | ❌ No | Users who also contribute | -| **Development** | `./setup-env.sh --uv` | ✅ Yes | Contributors, testing changes | - -#### Quick Start - -```bash -# For end users (global install, no activation needed) -./setup-env.sh --pipx -coldpress --version - -# For developers (local venv, requires activation) -./setup-env.sh --uv -source .venv/bin/activate -coldpress --version -``` - -#### For End Users (pipx) - -**pipx** installs Coldpress in an isolated environment with global CLI access - no activation needed. - -```bash -# Install with pipx (recommended for end users) -./setup-env.sh --pipx - -# Commands work globally, from any directory -coldpress --help -coldpress-setup --help -``` - -**Install in editable mode** (get updates as you pull from git): -```bash -./setup-env.sh --pipx-editable -``` - -**Manage installation:** -```bash -pipx upgrade coldpress # Upgrade to latest version -pipx reinstall coldpress # Reinstall -pipx uninstall coldpress # Remove completely -pipx list # Show installed packages -``` - -#### For Developers (uv) - -**uv** provides fast, reproducible virtual environments for development work. - -```bash -# One-time setup (installs uv, creates venv, installs coldpress) -./setup-env.sh --uv -# or just: -./setup-env.sh # --uv is the default - -# Activate the environment -source .venv/bin/activate - -# Use coldpress (while venv is active) -coldpress --help - -# Or use without activation -.venv/bin/coldpress --help -``` - -**For subsequent sessions:** -```bash -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 - -### 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: - -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) - -This is a one-time setup process that configures the cluster infrastructure for all users. - -### For Users - -Once the admin has completed the cluster setup, follow the **[User Quickstart Guide](docs/quickstart_user.md)** to: - -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) - -This workflow is repeatable for each job you want to run. - -## Documentation - -- **[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) - -## Quickstart Guides - -- **[Admin Quickstart](docs/quickstart_admin.md)** - Cluster setup for administrators (one-time) -- **[User Quickstart](docs/quickstart_user.md)** - Running workloads for users (repeatable) - - -## 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 -│ ├── 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 -``` - -## Testing - -Coldpress includes a comprehensive test suite that validates all fixes for [GitHub issue #37](https://github.com/asanaullah/coldpress/issues/37). - -### Run All Tests - -```bash -# From repository root -./tests/run_all_tests.sh -``` - -### Run Individual Tests - -```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 -``` - -### CI/CD - -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. - -See `tests/README.md` for detailed test documentation. - -## 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) - -**Local development:** -- Python 3.9+ (tested on Python 3.14) - -**Cluster tools:** -- `kubectl` or `oc` CLI (tested with oc 4.17.0) - -## Environment Variables - -Customize directory locations with environment variables: - -**For `coldpress` (job generation):** -- `COLDPRESS_DISCOVERY_DIR` - Discovery templates directory (default: `discovery`) -- `COLDPRESS_PROJECT_DIR` - Project configs directory (default: `projects`) -- `COLDPRESS_OUTPUT_DIR` - Default output directory (default: `output`) - -**For `coldpress-setup` (manifest generation):** -- `COLDPRESS_MANIFESTS_DIR` - Manifest output directory (default: `manifests`) -- `COLDPRESS_CLUSTER_DIR` - Cluster configs directory (default: `cluster`) -- `COLDPRESS_USER_DIR` - User configs directory (default: `users`) - -**Example:** -```bash -export COLDPRESS_OUTPUT_DIR=jobs -coldpress generate --config examples/pytorch_ddp_training/config.yaml -# Outputs to: jobs/ddp-training-job/ instead of output/ddp-training-job/ - -export COLDPRESS_MANIFESTS_DIR=gitops/manifests -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 - -output: ddp-training-job - -# Files to mount into container (creates ConfigMap) -files: - - train.py - - model_config.json -``` - -**Generate and run:** -```bash -coldpress generate --config examples/pytorch_ddp_training/config.yaml -cd output/ddp-training-job/ -./run.sh -``` - -**Results structure in PVC:** -``` -/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: - -```bash -coldpress generate --config examples/vllm_guidellm_benchmark/config.yaml -cd output/vllm-benchmark-job/ -./run.sh -``` - -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..b12f6d7 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 @@ -76,7 +76,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 +92,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/cli.py b/coldpress/cli.py index d790c8e..f41b837 100644 --- a/coldpress/cli.py +++ b/coldpress/cli.py @@ -1,5 +1,5 @@ # Generated by: Claude Sonnet 4.5 -"""Coldpress CLI - Generate JobSet manifests and scripts for AI/HPC workloads.""" +"""Coldpress CLI - Generate JobSet and Kubeflow manifests for AI/HPC workloads.""" import click import yaml @@ -7,10 +7,16 @@ import json 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 kubeflow_to_yaml from .script_gen import write_scripts from coldpress_common import ( validate_config, + validate_intent, validate_project_config, validate_task_specs, ) @@ -70,8 +76,56 @@ def _load_and_validate_config( return config_data, config_dir, project, discovery, output_name, nodes_from_config +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 intent schema + try: + validated_intent = validate_intent(intent_data) + except ValidationError as e: + raise ValueError(f"Intent validation failed: {e}") from e + + return validated_intent + + +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 at {job_spec_path}") + + with open(job_spec_path, "r") as f: + manifests = list(yaml.safe_load_all(f)) + + if not manifests: + raise ValueError("job-spec.yaml is empty") + + # Extract Jobs and pass through other resources + jobs = {} + other_resources = [] + + 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") + + return jobs, other_resources + + def _load_task_specs(config_dir): - """Load and validate task specifications from job-spec.yaml.""" + """Load and validate task specifications from job-spec.yaml (legacy format).""" job_spec_path = os.path.join(config_dir, "job-spec.yaml") if not os.path.exists(job_spec_path): raise FileNotFoundError(f"Could not find job-spec.yaml in {config_dir}") @@ -164,23 +218,41 @@ def _allocate_nodes_for_tasks(task_specs, manual_nodes): def _write_output_files( output_dir, - jobset, + manifest, services, job_spec, base_dir, config_dir, task_specs, node_assignments, + manifest_type="jobset", ): """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 + # Write manifest YAML based on type + if manifest_type == "jobset": + manifest_file = os.path.join(output_dir, "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_dir, "pytorchjob.yaml") + with open(manifest_file, "w") as f: + f.write(kubeflow_to_yaml(manifest)) + elif manifest_type == "tfjob": + manifest_file = os.path.join(output_dir, "tfjob.yaml") + with open(manifest_file, "w") as f: + f.write(kubeflow_to_yaml(manifest)) + elif manifest_type == "mpijob": + manifest_file = os.path.join(output_dir, "mpijob.yaml") + with open(manifest_file, "w") as f: + f.write(kubeflow_to_yaml(manifest)) + elif manifest_type == "inferenceservice": + manifest_file = os.path.join(output_dir, "inferenceservice.yaml") + with open(manifest_file, "w") as f: + f.write(kubeflow_to_yaml(manifest)) + + # Write Services YAML if any (JobSet only) if services: services_file = os.path.join(output_dir, "services.yaml") with open(services_file, "w") as f: @@ -204,6 +276,7 @@ def _write_output_files( metadata = { "job_name": job_spec["name"], "namespace": job_spec["namespace"], + "backend": manifest_type, "generated_at": datetime.now(timezone.utc).isoformat(), "base_dir": base_dir, "node_assignments": node_assignments, @@ -222,12 +295,13 @@ def _write_output_files( ) write_scripts( output_dir, - f"coldpress-{job_spec['name']}", + job_spec["name"], # Use original job name, not prefixed job_spec["namespace"], storage_pvc, base_dir, configmap_name, configmap_files, + manifest_type, # Pass manifest type to scripts ) return output_dir @@ -235,45 +309,22 @@ def _write_output_files( @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 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 intent file specifies: + - Target backend (jobset or kubeflow) + - Transformations to apply to vanilla k8s job-spec.yaml + - Task replication and dependencies + + 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 +332,201 @@ 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) + # Load and validate intent.yaml + intent_config = _load_intent(intent) + intent_dir = os.path.dirname(os.path.abspath(intent)) - project_config, namespace = _load_project_config(project) + # 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) - # Build job spec - job_spec = { - "name": job_name, - "tasks": task_specs, - "namespace": namespace, - "storage": project_config.get("storage", {}), - } - - # 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"], - } + # Load project config + project_config, namespace = _load_project_config(intent_config.project) - # 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 + # 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_map = { + "kubeflow": "Kubeflow", + "kserve": "KServe", + "jobset": "JobSet", + } + 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) + from .kubeflow_generator import ( + generate_pytorchjob_from_intent, + kubeflow_to_yaml, ) - 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, base_dir = generate_pytorchjob_from_intent( + vk8s_jobs, intent_config, project_config, namespace ) - else: - # Default: let Kubernetes scheduler decide - click.echo( - "Node scheduling: Kubernetes will select any coldpress-labeled node" + manifest_type = "pytorchjob" + services = [] # Kubeflow doesn't use separate services + elif backend == "kserve": + # Generate KServe manifest (InferenceService) + from .kubeflow_generator import ( + generate_inferenceservice_from_intent, + kubeflow_to_yaml, ) - 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") + manifest, base_dir = generate_inferenceservice_from_intent( + vk8s_jobs, intent_config, project_config, namespace + ) + manifest_type = "inferenceservice" + services = [] # KServe doesn't use separate services + else: + # Generate JobSet manifest + manifest, services, base_dir = generate_jobset_from_intent( + vk8s_jobs, intent_config, project_config, namespace + ) + manifest_type = "jobset" + + # Write manifests + os.makedirs(output_path, exist_ok=True) + + 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)) + + 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: + import shutil + + 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 + 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", {}) ) - 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) + volumes = job_template.get("spec", {}).get("volumes", []) + for volume in volumes: + if "configMap" in volume: + configmap_name = volume["configMap"]["name"] + break + if configmap_name: + break + elif manifest_type == "pytorchjob": + # Search for configMap volumes in PyTorchJob manifest + replica_specs = manifest.get("spec", {}).get("pytorchReplicaSpecs", {}) + for replica_type, replica_spec in replica_specs.items(): + volumes = ( + replica_spec.get("template", {}) + .get("spec", {}) + .get("volumes", []) + ) + for volume in volumes: + if "configMap" in volume: + configmap_name = volume["configMap"]["name"] + break + if configmap_name: + break + + # 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 = project_config.get("storage", {}).get( + "results", f"coldpress-{namespace}-storage" + ) + # 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 all output files - output_path = _write_output_files( + 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") diff --git a/coldpress/constants.py b/coldpress/constants.py new file mode 100644 index 0000000..533aa0d --- /dev/null +++ b/coldpress/constants.py @@ -0,0 +1,39 @@ +# Generated by: Claude Sonnet 4.5 +"""Coldpress constants - centralized configuration values.""" + +# Version +COLDPRESS_VERSION = "0.2.0" + +# Resource naming +COLDPRESS_PREFIX = "coldpress" +COLDPRESS_LABEL_MANAGED_BY = "coldpress" + +# Images +MKDIR_IMAGE = "registry.access.redhat.com/ubi9/ubi-minimal:latest" + +# Labels +COLDPRESS_LABELS = { + "app.kubernetes.io/managed-by": COLDPRESS_LABEL_MANAGED_BY, + "app.kubernetes.io/version": COLDPRESS_VERSION, +} + + +# 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" diff --git a/coldpress/generator.py b/coldpress/generator.py index 2ea616b..ab0da6d 100644 --- a/coldpress/generator.py +++ b/coldpress/generator.py @@ -1,12 +1,7 @@ # Generated by: Claude Sonnet 4.5 -"""JobSet YAML generation for Coldpress jobs.""" +"""Common generator utilities for Coldpress.""" -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 @@ -60,257 +55,16 @@ def generate_base_dir(namespace, job_name): 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. +def extract_container_config(task): + """Extract container configuration from task. Args: - job_spec: Job specification dict with name, namespace, tasks, storage - node_assignments: Dict mapping task index to node ID + task: Task specification dict Returns: - tuple: (JobSet manifest, services list, base_dir) + dict: Container configuration with keys: image, command, args, env, + working_dir, resources, gpus """ - 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] @@ -321,6 +75,8 @@ def _extract_container_config(task): "env": container_def.get("env", []), "working_dir": container_def.get("workingDir"), "resources": container_def.get("resources", {}), + "ports": container_def.get("ports", []), + "readinessProbe": container_def.get("readinessProbe"), "gpus": int( container_def.get("resources", {}) .get("requests", {}) @@ -336,540 +92,125 @@ def _extract_container_config(task): "env": task.get("env", []), "working_dir": task.get("workingDir"), "resources": task.get("resources", {}), + "ports": [], + "readinessProbe": None, "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_volumes_from_task(task, job_spec): + """Build volume list from task and job spec. -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) + Args: + task: Task specification dict + job_spec: Job specification dict - container = { - "name": "main", - "image": config["image"], - "volumeMounts": _build_volume_mounts(task, task_id, base_dir), - "resources": _build_container_resources(config["resources"], config["gpus"]), - } + Returns: + list: List of volume dicts for pod spec + """ + volumes = [] - 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", []): + # Add volumes from task + task_volumes = task.get("volumes", []) + for vol in task_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", - } - ] - } - ] + if vol_type == "pvc": + # Use storage config from job_spec + storage = job_spec.get("storage", {}) + pvc_name = storage.get( + vol_name, f"coldpress-{job_spec['namespace']}-{vol_name}" + ) + volumes.append( + {"name": vol_name, "persistentVolumeClaim": {"claimName": pvc_name}} + ) + elif vol_type == "emptyDir": + volumes.append( + { + "name": vol_name, + "emptyDir": { + "medium": vol.get("medium", ""), + "sizeLimit": vol.get("sizeLimit", ""), + }, } - } - } - - _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 + return volumes -def build_discovery_init_container(template_path, task_id, base_dir, data_pvc_name): - """ - Build discovery init container from template. +def build_volume_mounts_from_task(task): + """Build volume mount list from task. 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 + task: Task specification dict Returns: - dict: Init container spec, or None if template not found + list: List of volume mount dicts """ - 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) + volume_mounts = [] - 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" + task_volumes = task.get("volumes", []) + for vol in task_volumes: + if "mount" in vol: + volume_mounts.append({"name": vol["name"], "mountPath": vol["mount"]}) - 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}, - } - ], - }, - }, - } - }, - } + return volume_mounts -def build_discovery_job(template_path, base_dir, data_pvc_name, node_id): - """ - Build discovery job from template. +def detect_job_type(task): + """Detect job type from task configuration. 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 + task: Task specification dict Returns: - dict: ReplicatedJob for discovery, or None if template not found + str: 'training' or 'inference' """ - 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"}], - } + container_config = extract_container_config(task) - # 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", - } - ] - } - ] - } - } - } + # Check for readinessProbe - typical of inference services + if container_config.get("readinessProbe"): + return "inference" - 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 + # Check image for inference patterns + image = container_config.get("image", "").lower() + if any( + pattern in image + for pattern in ["serving", "inference", "kserve", "torchserve", "triton"] + ): + return "inference" + + # Default to training + return "training" -def jobset_to_yaml(jobset): - """Convert JobSet dict to YAML string.""" - return yaml.dump(jobset, default_flow_style=False, sort_keys=False) +def detect_framework(task): + """Detect ML framework from task configuration. + Args: + task: Task specification dict -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) + Returns: + str: 'pytorch', 'tensorflow', 'mpi', or None + """ + container_config = extract_container_config(task) + image = container_config.get("image", "").lower() + command = container_config.get("command", []) + + # Check image name + if "pytorch" in image: + return "pytorch" + elif "tensorflow" in image or "tf" in image: + return "tensorflow" + + # Check command + command_str = " ".join(command) if isinstance(command, list) else str(command) + if "torch" in command_str.lower(): + return "pytorch" + elif "tensorflow" in command_str.lower(): + return "tensorflow" + elif "mpirun" in command_str.lower(): + return "mpi" + + return None diff --git a/coldpress/jobset_generator.py b/coldpress/jobset_generator.py new file mode 100644 index 0000000..55c246d --- /dev/null +++ b/coldpress/jobset_generator.py @@ -0,0 +1,678 @@ +# Generated 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 os +import re +import copy +import sys +import yaml +from datetime import datetime, timezone +from .constants import ( + COLDPRESS_LABELS, + MKDIR_IMAGE, + get_kueue_queue_label, + get_jobset_name, + get_pvc_name, +) + + +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 + dangerous_chars = ["`", "$", "|", ";", "&", ">", "<", "\n", "\r"] + for char in dangerous_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 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 replace_or_add_arg(args: list, key: str, value: str) -> list: + """Replace or add argument in args list.""" + new_args = [] + i = 0 + found = False + + while i < len(args): + arg = str(args[i]) + + if arg.startswith(f"--{key}="): + new_args.append(f"--{key}={value}") + found = True + i += 1 + elif arg == f"--{key}": + new_args.append(arg) + new_args.append(value) + found = True + i += 2 + else: + new_args.append(arg) + i += 1 + + if not found: + new_args.append(f"--{key}={value}") + + return new_args + + +def build_discovery_init_container( + template_path: str, task_id: int, base_dir: str, storage_volume_name: str +): + """ + 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 + storage_volume_name: Name of the storage volume to mount + + 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 = copy.deepcopy(containers[0]) + 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": storage_volume_name, + "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 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) + storage = project_config.get("storage", {}) + data_pvc_name = storage.get("results", get_pvc_name(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 = 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", "discovery") + 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) + + # 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.""" + # 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 replacements + if task_intent.args and container.get("args"): + args = ( + container["args"] + if isinstance(container["args"], list) + else [container["args"]] + ) + + for arg_key, arg_value in task_intent.args.items(): + # Validate arg key follows Kubernetes conventions + if ( + not arg_key.replace("_", "-") + .replace(".", "-") + .replace("/", "-") + .isalnum() + ): + # Allow alphanumeric, _, -, ., / in arg keys + pass # Most arg keys are valid + + # Substitute macros with validation + context = f" for arg '--{arg_key}'" + substituted_value = substitute_macros(arg_value, macros, context) + args = replace_or_add_arg(args, arg_key, substituted_value) + + container["args"] = args + + # Apply env var additions/replacements + if task_intent.env: + # 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 + + # 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, task_counter, base_dir, storage_volume_name + ) + 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.""" + import yaml + + return yaml.dump(jobset, default_flow_style=False, sort_keys=False) + + +def services_to_yaml(services: list) -> str: + """Convert Services list to YAML string.""" + import yaml + + return yaml.dump_all(services, default_flow_style=False, sort_keys=False) diff --git a/coldpress/jobset_generator_legacy.py b/coldpress/jobset_generator_legacy.py new file mode 100644 index 0000000..790bc0f --- /dev/null +++ b/coldpress/jobset_generator_legacy.py @@ -0,0 +1,803 @@ +# Generated by: Claude Sonnet 4.5 +"""JobSet YAML generation for Coldpress jobs.""" + +import os +import sys +import yaml +import shlex +from urllib.parse import urlparse + +# Import common utilities +from .generator import ( + COLDPRESS_LABELS, + generate_base_dir, + extract_container_config, +) + + +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 + + +# Use extract_container_config from common generator.py +_extract_container_config = extract_container_config + + +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/kubeflow_generator.py b/coldpress/kubeflow_generator.py new file mode 100644 index 0000000..448b7a1 --- /dev/null +++ b/coldpress/kubeflow_generator.py @@ -0,0 +1,1076 @@ +# Generated by: Claude Sonnet 4.5 +"""Kubeflow manifest generation for Coldpress jobs.""" + +import yaml +from .generator import ( + COLDPRESS_LABELS, + generate_base_dir, + extract_container_config, + build_volumes_from_task, + build_volume_mounts_from_task, + detect_job_type, + detect_framework, +) + + +def generate_kubeflow(job_spec, node_assignments): + """Generate Kubeflow manifest (PyTorchJob or InferenceService). + + Auto-detects job type: + - Training job → PyTorchJob/TFJob/MPIJob + - Inference job → KServe InferenceService + + Args: + job_spec: Job specification dict + node_assignments: Node assignment dict + + Returns: + tuple: (manifest_dict, base_dir, manifest_type) + manifest_type is 'pytorchjob', 'tfjob', 'mpijob', or 'inferenceservice' + """ + tasks = job_spec["tasks"] + if len(tasks) == 0: + raise ValueError("No tasks specified in job spec") + + # For now, only handle single-task jobs + task = tasks[0] + job_type = detect_job_type(task) + + if job_type == "inference": + return generate_inferenceservice(job_spec, node_assignments) + else: + # Training job + framework = detect_framework(task) + if framework == "pytorch": + return generate_pytorchjob(job_spec, node_assignments) + elif framework == "tensorflow": + return generate_tfjob(job_spec, node_assignments) + elif framework == "mpi": + return generate_mpijob(job_spec, node_assignments) + else: + # Default to PyTorchJob if framework unclear + return generate_pytorchjob(job_spec, node_assignments) + + +def _build_mkdir_init_container(base_dir, storage_pvc_name): + """Build init container to create directory structure. + + Args: + base_dir: Base directory path to create + storage_pvc_name: Name of the PVC for storage + + Returns: + dict: Init container spec + """ + return { + "name": "mkdir", + "image": "registry.access.redhat.com/ubi9/ubi-minimal:latest", + "command": [ + "sh", + "-c", + f"mkdir -p /mnt/coldpress-data/{base_dir} && echo Created directory {base_dir}", + ], + "volumeMounts": [ + {"name": "coldpress-data", "mountPath": "/mnt/coldpress-data"} + ], + } + + +def generate_pytorchjob(job_spec, node_assignments): + """Generate PyTorchJob manifest from job spec. + + Args: + job_spec: Job specification dict + node_assignments: Node assignment dict + + Returns: + tuple: (pytorchjob_dict, base_dir, 'pytorchjob') + """ + namespace = job_spec["namespace"] + tasks = job_spec["tasks"] + task = tasks[0] # Single task for now + job_name = job_spec["name"] + + # Generate base directory for results + base_dir = generate_base_dir(namespace, job_name) + + # Get storage PVC name + storage = job_spec.get("storage", {}) + storage_pvc_name = storage.get("results", f"coldpress-{namespace}-storage") + + # Extract container configuration + container_config = extract_container_config(task) + + # Build container spec + container = { + "name": container_config.get("name", "pytorch"), + "image": container_config["image"], + } + + if container_config.get("command"): + container["command"] = container_config["command"] + if container_config.get("args"): + container["args"] = container_config["args"] + if container_config.get("working_dir"): + container["workingDir"] = container_config["working_dir"] + if container_config.get("env"): + container["env"] = container_config["env"] + + # Add resources + if container_config.get("resources"): + container["resources"] = container_config["resources"] + + # Build volume mounts + volume_mounts = build_volume_mounts_from_task(task) + + # Add ConfigMap volume mount if present + if "configmap" in job_spec: + volume_mounts.append({"name": "config-files", "mountPath": "/workspace"}) + + if volume_mounts: + container["volumeMounts"] = volume_mounts + + # Build volumes + volumes = build_volumes_from_task(task, job_spec) + + # Ensure coldpress-data PVC is in volumes (needed for init container) + if not any(v.get("name") == "coldpress-data" for v in volumes): + volumes.insert( + 0, + { + "name": "coldpress-data", + "persistentVolumeClaim": {"claimName": storage_pvc_name}, + }, + ) + + # Add ConfigMap volume if present + if "configmap" in job_spec: + cm_name = job_spec["configmap"]["name"] + # Add coldpress- prefix if not already present + if not cm_name.startswith("coldpress-"): + cm_name = f"coldpress-{cm_name}" + volumes.append({"name": "config-files", "configMap": {"name": cm_name}}) + + # Add init container to create directories + mkdir_init = _build_mkdir_init_container(base_dir, storage_pvc_name) + + # Build pod template + pod_template = { + "metadata": { + "annotations": { + "sidecar.istio.io/inject": "false" # Prescriptive: disable Istio + } + }, + "spec": { + "initContainers": [mkdir_init], + "containers": [container], + "restartPolicy": "OnFailure", + }, + } + + if volumes: + pod_template["spec"]["volumes"] = volumes + + # Apply node selector if specified + node_id = node_assignments.get(0, "any") + if node_id != "any": + pod_template["spec"]["nodeSelector"] = {"coldpress.node": str(node_id)} + + # Apply tolerations if task requires + if task.get("tolerate_all"): + pod_template["spec"]["tolerations"] = [{"operator": "Exists"}] + + # Build PyTorchJob manifest + pytorchjob_name = f"coldpress-{job_name}" + pytorchjob_labels = COLDPRESS_LABELS.copy() + pytorchjob_labels.update( + { + "coldpress.io/job-id": pytorchjob_name, + "kueue.x-k8s.io/queue-name": f"coldpress-local-queue-{namespace}", # Required for Kueue integration + } + ) + + # Add coldpress labels to pod template + if "labels" not in pod_template["metadata"]: + pod_template["metadata"]["labels"] = {} + pod_template["metadata"]["labels"].update(COLDPRESS_LABELS.copy()) + pod_template["metadata"]["labels"]["coldpress.io/job-id"] = pytorchjob_name + + pytorchjob = { + "apiVersion": "kubeflow.org/v1", + "kind": "PyTorchJob", + "metadata": { + "name": pytorchjob_name, + "namespace": namespace, + "labels": pytorchjob_labels, + "annotations": { + "coldpress.io/base-dir": base_dir, + "kueue.x-k8s.io/wait-for-pods-ready-timeout": "20m", # Allow time for image pulls + }, + }, + "spec": { + "pytorchReplicaSpecs": { + "Master": { + "replicas": 1, + "restartPolicy": "OnFailure", + "template": pod_template, + } + } + }, + } + + return pytorchjob, base_dir, "pytorchjob" + + +def generate_tfjob(job_spec, node_assignments): + """Generate TFJob manifest from job spec. + + Args: + job_spec: Job specification dict + node_assignments: Node assignment dict + + Returns: + tuple: (tfjob_dict, base_dir, 'tfjob') + """ + namespace = job_spec["namespace"] + tasks = job_spec["tasks"] + task = tasks[0] + job_name = job_spec["name"] + + base_dir = generate_base_dir(namespace, job_name) + + # Get storage PVC name + storage = job_spec.get("storage", {}) + storage_pvc_name = storage.get("results", f"coldpress-{namespace}-storage") + + container_config = extract_container_config(task) + + container = { + "name": "tensorflow", + "image": container_config["image"], + } + + if container_config.get("command"): + container["command"] = container_config["command"] + if container_config.get("args"): + container["args"] = container_config["args"] + if container_config.get("resources"): + container["resources"] = container_config["resources"] + + volume_mounts = build_volume_mounts_from_task(task) + if volume_mounts: + container["volumeMounts"] = volume_mounts + + volumes = build_volumes_from_task(task, job_spec) + + # Ensure coldpress-data PVC is in volumes (needed for init container) + if not any(v.get("name") == "coldpress-data" for v in volumes): + volumes.insert( + 0, + { + "name": "coldpress-data", + "persistentVolumeClaim": {"claimName": storage_pvc_name}, + }, + ) + + # Add init container to create directories + mkdir_init = _build_mkdir_init_container(base_dir, storage_pvc_name) + + pod_template = { + "metadata": {"annotations": {"sidecar.istio.io/inject": "false"}}, + "spec": { + "initContainers": [mkdir_init], + "containers": [container], + "restartPolicy": "OnFailure", + }, + } + + if volumes: + pod_template["spec"]["volumes"] = volumes + + # Add coldpress labels to pod template + if "labels" not in pod_template["metadata"]: + pod_template["metadata"]["labels"] = {} + tfjob_name = f"coldpress-{job_name}" + pod_template["metadata"]["labels"].update(COLDPRESS_LABELS.copy()) + pod_template["metadata"]["labels"]["coldpress.io/job-id"] = tfjob_name + + tfjob_labels = COLDPRESS_LABELS.copy() + tfjob_labels.update( + { + "coldpress.io/job-id": tfjob_name, + "kueue.x-k8s.io/queue-name": f"coldpress-local-queue-{namespace}", + } + ) + + tfjob = { + "apiVersion": "kubeflow.org/v1", + "kind": "TFJob", + "metadata": { + "name": tfjob_name, + "namespace": namespace, + "labels": tfjob_labels, + "annotations": { + "coldpress.io/base-dir": base_dir, + "kueue.x-k8s.io/wait-for-pods-ready-timeout": "20m", # Allow time for image pulls + }, + }, + "spec": { + "tfReplicaSpecs": { + "Worker": { + "replicas": 1, + "restartPolicy": "OnFailure", + "template": pod_template, + } + } + }, + } + + return tfjob, base_dir, "tfjob" + + +def generate_mpijob(job_spec, node_assignments): + """Generate MPIJob manifest from job spec. + + Args: + job_spec: Job specification dict + node_assignments: Node assignment dict + + Returns: + tuple: (mpijob_dict, base_dir, 'mpijob') + """ + namespace = job_spec["namespace"] + tasks = job_spec["tasks"] + task = tasks[0] + job_name = job_spec["name"] + + base_dir = generate_base_dir(namespace, job_name) + + # Get storage PVC name + storage = job_spec.get("storage", {}) + storage_pvc_name = storage.get("results", f"coldpress-{namespace}-storage") + + container_config = extract_container_config(task) + + launcher_container = { + "name": "mpi-launcher", + "image": container_config["image"], + } + + if container_config.get("command"): + launcher_container["command"] = container_config["command"] + if container_config.get("args"): + launcher_container["args"] = container_config["args"] + if container_config.get("resources"): + launcher_container["resources"] = container_config["resources"] + + worker_container = launcher_container.copy() + worker_container["name"] = "mpi-worker" + + mpijob_name = f"coldpress-{job_name}" + mpijob_labels = COLDPRESS_LABELS.copy() + mpijob_labels.update( + { + "coldpress.io/job-id": mpijob_name, + "kueue.x-k8s.io/queue-name": f"coldpress-local-queue-{namespace}", + } + ) + + # Add init container to create directories + mkdir_init = _build_mkdir_init_container(base_dir, storage_pvc_name) + + # Build pod templates with coldpress labels + launcher_pod_template = { + "metadata": { + "labels": COLDPRESS_LABELS.copy(), + "annotations": {"sidecar.istio.io/inject": "false"}, + }, + "spec": {"initContainers": [mkdir_init], "containers": [launcher_container]}, + } + launcher_pod_template["metadata"]["labels"]["coldpress.io/job-id"] = mpijob_name + + worker_pod_template = { + "metadata": { + "labels": COLDPRESS_LABELS.copy(), + "annotations": {"sidecar.istio.io/inject": "false"}, + }, + "spec": {"initContainers": [mkdir_init], "containers": [worker_container]}, + } + worker_pod_template["metadata"]["labels"]["coldpress.io/job-id"] = mpijob_name + + mpijob = { + "apiVersion": "kubeflow.org/v2beta1", + "kind": "MPIJob", + "metadata": { + "name": mpijob_name, + "namespace": namespace, + "labels": mpijob_labels, + "annotations": { + "coldpress.io/base-dir": base_dir, + "kueue.x-k8s.io/wait-for-pods-ready-timeout": "20m", # Allow time for image pulls + }, + }, + "spec": { + "slotsPerWorker": 1, + "runPolicy": {"cleanPodPolicy": "Running"}, + "mpiReplicaSpecs": { + "Launcher": {"replicas": 1, "template": launcher_pod_template}, + "Worker": {"replicas": 1, "template": worker_pod_template}, + }, + }, + } + + return mpijob, base_dir, "mpijob" + + +def generate_inferenceservice(job_spec, node_assignments): + """Generate KServe InferenceService manifest from job spec. + + Args: + job_spec: Job specification dict + node_assignments: Node assignment dict + + Returns: + tuple: (inferenceservice_dict, base_dir, 'inferenceservice') + """ + namespace = job_spec["namespace"] + tasks = job_spec["tasks"] + task = tasks[0] + job_name = job_spec["name"] + + base_dir = generate_base_dir(namespace, job_name) + + container_config = extract_container_config(task) + framework = detect_framework(task) + + # Build predictor based on framework + predictor = {} + + if framework == "pytorch": + predictor["pytorch"] = { + "storageUri": task.get("storageUri", "gs://bucket/model"), + "resources": container_config.get("resources", {}), + } + elif framework == "tensorflow": + predictor["tensorflow"] = { + "storageUri": task.get("storageUri", "gs://bucket/model"), + "resources": container_config.get("resources", {}), + } + else: + # Use custom container + predictor["containers"] = [ + { + "name": "kserve-container", + "image": container_config["image"], + "resources": container_config.get("resources", {}), + } + ] + + inferenceservice_name = f"coldpress-{job_name}" + inferenceservice_labels = COLDPRESS_LABELS.copy() + inferenceservice_labels.update( + { + "coldpress.io/job-id": inferenceservice_name, + } + ) + + inferenceservice = { + "apiVersion": "serving.kserve.io/v1beta1", + "kind": "InferenceService", + "metadata": { + "name": inferenceservice_name, + "namespace": namespace, + "labels": inferenceservice_labels, + "annotations": { + "coldpress.io/base-dir": base_dir, + }, + }, + "spec": {"predictor": predictor}, + } + + return inferenceservice, base_dir, "inferenceservice" + + +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) + """ + import copy + import re + from datetime import datetime, timezone + from .constants import ( + COLDPRESS_LABELS, + get_kueue_queue_label, + ) + + # 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}" + + storage = project_config.get("storage", {}) + data_pvc_name = storage.get("results", f"coldpress-{namespace}-storage") + + # 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 + + # Build macros for arg/env substitution + def substitute_macros(value: str, macros: dict) -> str: + """Substitute ${MACRO} patterns in a string.""" + if not isinstance(value, str): + return value + + pattern = r"\$\{([^}]+)\}" + + def replacer(match): + macro_name = match.group(1) + if macro_name not in macros: + return match.group(0) + return str(macros[macro_name]) + + return re.sub(pattern, replacer, value) + + # KServe macros + macros = { + "REPLICAS": str(replicas), + "TASK_NAME": task_name, + } + + # Apply arg replacements + if task_intent.args and container.get("args"): + args = ( + container["args"] + if isinstance(container["args"], list) + else [container["args"]] + ) + + for arg_key, arg_value in task_intent.args.items(): + substituted_value = substitute_macros(arg_value, macros) + + # Replace or add argument + new_args = [] + i = 0 + found = False + + while i < len(args): + arg = str(args[i]) + + if arg.startswith(f"--{arg_key}="): + new_args.append(f"--{arg_key}={substituted_value}") + found = True + i += 1 + elif arg == f"--{arg_key}": + new_args.append(arg) + new_args.append(substituted_value) + found = True + i += 2 + else: + new_args.append(arg) + i += 1 + + if not found: + new_args.append(f"--{arg_key}={substituted_value}") + + args = new_args + + container["args"] = args + + # Apply env var additions/replacements + if task_intent.env: + env_list = container.get("env", []) + if not isinstance(env_list, list): + env_list = [] + + for env_key, env_value in task_intent.env.items(): + substituted_value = substitute_macros(str(env_value), macros) + + # 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 + + # 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.dump(manifest, default_flow_style=False, sort_keys=False) + + +def _build_discovery_init_container( + template_path: str, base_dir: str, storage_volume_name: str +): + """ + Build discovery init container from template for Kubeflow. + + Args: + template_path: Path to discovery template YAML + base_dir: Base directory path + storage_volume_name: Name of the storage volume to mount + + Returns: + dict: Init container spec, or None if template not found + """ + import os + import sys + import yaml + import copy + + 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 = copy.deepcopy(containers[0]) + container["name"] = "discovery" + + # Update volume mounts to use PVC with base directory + container["volumeMounts"] = [ + { + "name": storage_volume_name, + "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] + + 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 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) + """ + import copy + import os + import re + from datetime import datetime, timezone + from .constants import ( + COLDPRESS_LABELS, + MKDIR_IMAGE, + get_kueue_queue_label, + ) + + # 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}" + + storage = project_config.get("storage", {}) + data_pvc_name = storage.get("results", f"coldpress-{namespace}-storage") + + # 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 = 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", "discovery") + 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_task_names = set(discovery_tasks) + + # 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 + + # Build macros for arg substitution (using Kubeflow auto-injected vars) + def substitute_macros(value: str, macros: dict) -> str: + """Substitute ${MACRO} patterns in a string.""" + if not isinstance(value, str): + return value + + pattern = r"\$\{([^}]+)\}" + + def replacer(match): + macro_name = match.group(1) + if macro_name not in macros: + return match.group(0) + return str(macros[macro_name]) + + return re.sub(pattern, replacer, value) + + # Kubeflow macros - simpler because PyTorchJob auto-injects env vars + macros = { + "REPLICAS": str(replicas), + "TASK_NAME": task_name, + } + + # Apply arg replacements + if task_intent.args and container.get("args"): + args = ( + container["args"] + if isinstance(container["args"], list) + else [container["args"]] + ) + + for arg_key, arg_value in task_intent.args.items(): + substituted_value = substitute_macros(arg_value, macros) + + # Replace or add argument + new_args = [] + i = 0 + found = False + + while i < len(args): + arg = str(args[i]) + + if arg.startswith(f"--{arg_key}="): + new_args.append(f"--{arg_key}={substituted_value}") + found = True + i += 1 + elif arg == f"--{arg_key}": + new_args.append(arg) + new_args.append(substituted_value) + found = True + i += 2 + else: + new_args.append(arg) + i += 1 + + if not found: + new_args.append(f"--{arg_key}={substituted_value}") + + args = new_args + + container["args"] = args + + # Apply env var additions/replacements + if task_intent.env: + env_list = container.get("env", []) + if not isinstance(env_list, list): + env_list = [] + + for env_key, env_value in task_intent.env.items(): + substituted_value = substitute_macros(str(env_value), macros) + + # 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 + + # 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 = [] + + # 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 + ) + 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/script_gen.py b/coldpress/script_gen.py index fae459b..9d87b62 100644 --- a/coldpress/script_gen.py +++ b/coldpress/script_gen.py @@ -5,17 +5,23 @@ def generate_run_script( - job_name, namespace, timeout="1h", configmap_name=None, configmap_files=None + job_name, + namespace, + timeout="1h", + 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) 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 @@ -23,12 +29,74 @@ def generate_run_script( configmap_apply = "" if configmap_name and configmap_files: from_files = " ".join([f"--from-file={f}" for f in configmap_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 - """ + # Generate manifest-specific commands + if manifest_type == "jobset": + manifest_file = "jobset.yaml" + resource_type = "jobset" + resource_name = f"coldpress-{job_name}" + apply_cmd = 'echo "Applying JobSet..."\noc apply -f jobset.yaml' + services_apply = """ +# Apply services if they exist +if [ -f services.yaml ]; then + echo "Applying Services..." + oc apply -f services.yaml +fi""" + elif manifest_type == "pytorchjob": + manifest_file = "pytorchjob.yaml" + resource_type = "pytorchjob" + resource_name = f"coldpress-{job_name}" + apply_cmd = 'echo "Applying PyTorchJob..."\noc apply -f pytorchjob.yaml' + services_apply = "" + elif manifest_type == "tfjob": + manifest_file = "tfjob.yaml" + resource_type = "tfjob" + resource_name = f"coldpress-{job_name}" + apply_cmd = 'echo "Applying TFJob..."\noc apply -f tfjob.yaml' + services_apply = "" + elif manifest_type == "mpijob": + manifest_file = "mpijob.yaml" + resource_type = "mpijob" + resource_name = f"coldpress-{job_name}" + apply_cmd = 'echo "Applying MPIJob..."\noc apply -f mpijob.yaml' + services_apply = "" + elif manifest_type == "inferenceservice": + manifest_file = "kservejob.yaml" + resource_type = "inferenceservice" + resource_name = f"coldpress-{job_name}" + apply_cmd = ( + 'echo "Applying KServe InferenceService..."\noc apply -f kservejob.yaml' + ) + services_apply = "" + else: + manifest_file = "manifest.yaml" + resource_type = manifest_type + resource_name = f"coldpress-{job_name}" + apply_cmd = f'echo "Applying {manifest_type}..."\noc apply -f {manifest_file}' + services_apply = "" + + # 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,54 +106,87 @@ 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 """ 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 +197,76 @@ 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 """ + # 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 +276,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 +284,33 @@ def generate_logs_script(job_name, namespace, storage_pvc, base_dir): Returns: str: Bash script content """ + 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() {{ @@ -256,7 +387,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 +395,18 @@ def generate_explore_script(job_name, namespace, storage_pvc, base_dir): Returns: str: Bash script content """ + 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,7 +422,7 @@ 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: @@ -350,7 +483,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 +491,18 @@ def generate_copy_script(job_name, namespace, storage_pvc, base_dir): Returns: str: Bash script content """ + 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,7 +523,7 @@ 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: @@ -442,18 +577,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 +600,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_common/__init__.py b/coldpress_common/__init__.py index 2c488d0..491deec 100644 --- a/coldpress_common/__init__.py +++ b/coldpress_common/__init__.py @@ -3,6 +3,7 @@ from .model import ( validate_config, + validate_intent, validate_project_config, validate_task_specs, validate_job_spec, @@ -10,6 +11,9 @@ validate_cluster_config, validate_kubernetes_name, WorkloadConfig, + IntentConfig, + TaskIntent, + DependsOnConfig, ProjectConfig, TaskSpec, JobSpec, @@ -19,6 +23,7 @@ __all__ = [ "validate_config", + "validate_intent", "validate_project_config", "validate_task_specs", "validate_job_spec", @@ -26,6 +31,9 @@ "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..d4d71cc 100644 --- a/coldpress_common/model.py +++ b/coldpress_common/model.py @@ -242,8 +242,105 @@ def parse_tasks(cls, v): return v +class DependsOnConfig(BaseModel): + """Task dependency configuration.""" + + task: str + wait_for: Literal["ready", "completion"] + + +class TaskIntent(BaseModel): + """Task intent specification from intent.yaml.""" + + name: str + replicas: Optional[int] = 1 + nodes: Optional[list[int]] = None + args: Optional[dict[str, str]] = 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"] = "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 +361,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 +378,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 @@ -479,3 +577,19 @@ def validate_cluster_config(cluster_data: dict) -> ClusterConfig: ValidationError: If validation fails """ return ClusterConfig.model_validate(cluster_data) + + +def validate_intent(intent_data: dict) -> IntentConfig: + """ + Validate intent.yaml data. + + Args: + intent_data: Raw YAML data from intent file + + Returns: + Validated IntentConfig + + Raises: + ValidationError: If validation fails + """ + return IntentConfig.model_validate(intent_data) diff --git a/coldpress_setup/cli.py b/coldpress_setup/cli.py index 696c21e..928f5f4 100644 --- a/coldpress_setup/cli.py +++ b/coldpress_setup/cli.py @@ -121,7 +121,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 +158,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() @@ -309,7 +316,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 +325,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..3d718b0 100644 --- a/coldpress_setup/generator.py +++ b/coldpress_setup/generator.py @@ -143,8 +143,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 +164,45 @@ 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 + + # 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"], + } + ) + + # Common permissions (always needed) + rules.extend( + [ { "apiGroups": ["batch"], "resources": ["jobs"], @@ -193,9 +222,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) @@ -368,7 +417,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 +429,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 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/comp/job-spec.yaml b/examples/pytorch_ddp_training/comp/job-spec.yaml new file mode 100644 index 0000000..7d3229c --- /dev/null +++ b/examples/pytorch_ddp_training/comp/job-spec.yaml @@ -0,0 +1,75 @@ +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: ["bash", "-c"] + args: + - | + # Build torch.distributed.run command from env vars + CMD="python -m torch.distributed.run" + + # Add distributed params if env vars are set + [ -n "$NPROC_PER_NODE" ] && CMD="$CMD --nproc_per_node=$NPROC_PER_NODE" + [ -n "$NNODES" ] && CMD="$CMD --nnodes=$NNODES" + [ -n "$RANK" ] && CMD="$CMD --node_rank=$RANK" + [ -n "$MASTER_ADDR" ] && CMD="$CMD --master_addr=$MASTER_ADDR" + [ -n "$MASTER_PORT" ] && CMD="$CMD --master_port=$MASTER_PORT" + + # Training script and args + CMD="$CMD train.py" + CMD="$CMD --dataset=mnist" + CMD="$CMD --train-test-split=0.8" + CMD="$CMD --epochs=50" + CMD="$CMD --batch-size=128" + CMD="$CMD --hidden-size=4096" + CMD="$CMD --lr=0.01" + CMD="$CMD --output-dir=/results/checkpoints" + + echo "Running: $CMD" + exec $CMD + resources: + requests: + nvidia.com/gpu: "2" + memory: "16Gi" + cpu: "8" + limits: + nvidia.com/gpu: "2" + memory: "16Gi" + env: + - name: NCCL_DEBUG + value: "INFO" + - name: NPROC_PER_NODE + value: "2" + 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/comp/jobset.yaml b/examples/pytorch_ddp_training/comp/jobset.yaml new file mode 100644 index 0000000..ce6702a --- /dev/null +++ b/examples/pytorch_ddp_training/comp/jobset.yaml @@ -0,0 +1,596 @@ +apiVersion: jobset.x-k8s.io/v1alpha2 +kind: JobSet +metadata: + name: coldpress-ddp-training + namespace: coldpress-project + labels: + app.kubernetes.io/managed-by: coldpress + app.kubernetes.io/version: 0.2.0 + kueue.x-k8s.io/queue-name: coldpress-local-queue-coldpress-project + coldpress.io/job-id: coldpress-ddp-training + annotations: + coldpress.io/base-dir: coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821 + coldpress.io/storage-pvc: coldpress-project-storage +spec: + suspend: true + replicatedJobs: + - 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 + args: + - mkdir -p /data/coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821 + /data/coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821/task-0 + /data/coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821/task-1 + && echo Created directory /data/coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821 + with 2 task subdirectories + volumeMounts: + - name: storage + mountPath: /data + volumes: + - name: storage + persistentVolumeClaim: + claimName: coldpress-project-storage + - name: task-0 + replicas: 1 + template: + spec: + parallelism: 1 + completions: 1 + backoffLimit: 0 + template: + metadata: + labels: + app: task-0 + coldpress/gid: coldpress-ddp-training + spec: + restartPolicy: Never + containers: + - name: training + image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime + workingDir: /workspace + command: + - bash + - -c + args: + - '# Build torch.distributed.run command from env vars + + CMD="python -m torch.distributed.run" + + + # Add distributed params if env vars are set + + [ -n "$NPROC_PER_NODE" ] && CMD="$CMD --nproc_per_node=$NPROC_PER_NODE" + + [ -n "$NNODES" ] && CMD="$CMD --nnodes=$NNODES" + + [ -n "$RANK" ] && CMD="$CMD --node_rank=$RANK" + + [ -n "$MASTER_ADDR" ] && CMD="$CMD --master_addr=$MASTER_ADDR" + + [ -n "$MASTER_PORT" ] && CMD="$CMD --master_port=$MASTER_PORT" + + + # Training script and args + + CMD="$CMD train.py" + + CMD="$CMD --dataset=mnist" + + CMD="$CMD --train-test-split=0.8" + + CMD="$CMD --epochs=50" + + CMD="$CMD --batch-size=128" + + CMD="$CMD --hidden-size=4096" + + CMD="$CMD --lr=0.01" + + CMD="$CMD --output-dir=/results/checkpoints" + + + echo "Running: $CMD" + + exec $CMD + + ' + resources: + requests: + nvidia.com/gpu: '2' + memory: 16Gi + cpu: '8' + limits: + nvidia.com/gpu: '2' + memory: 16Gi + env: + - name: NCCL_DEBUG + value: INFO + - name: NPROC_PER_NODE + value: '2' + - name: NNODES + value: '2' + - name: MASTER_ADDR + value: coldpress-ddp-training-task-0-0-0.coldpress-ddp-training.coldpress-project.svc.cluster.local + - name: MASTER_PORT + value: '29500' + - name: RANK + value: '0' + - name: WORLD_SIZE + value: '2' + volumeMounts: + - name: results + mountPath: /results + subPath: coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821/task-0 + - name: dshm + mountPath: /dev/shm + - name: training-script + mountPath: /workspace/train.py + subPath: train.py + tolerations: &id001 + - operator: Exists + volumes: + - name: results + persistentVolumeClaim: + claimName: coldpress-project-storage + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 16Gi + - name: training-script + configMap: + name: coldpress-ddp-training-files + initContainers: + - name: discovery + image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime + command: + - /bin/bash + - -c + args: + - "echo \"User System Snapshot - Capturing performance baseline with\ + \ dynamic benchmarks...\"\n\n# Create comprehensive Python script\ + \ in /tmp\ncat << 'EOF_SNAPSHOT' > /tmp/generate_snapshot.py\nimport\ + \ subprocess\nimport json\nimport time\nimport os\nfrom datetime import\ + \ datetime\n\ndef run_cmd(command):\n try:\n out = subprocess.run(command,\ + \ shell=True, capture_output=True, text=True, timeout=10)\n \ + \ return out.stdout.strip() if out.returncode == 0 else \"Not available\"\ + \n except Exception as e:\n return \"Not available\"\n\n\ + def cpu_benchmark():\n \"\"\"Simple CPU compute benchmark - measures\ + \ MFLOPS\"\"\"\n start = time.time()\n result = 0.0\n iterations\ + \ = 10000000\n for i in range(iterations):\n result += i\ + \ * 1.5 - 0.7\n elapsed = time.time() - start\n mflops = (iterations\ + \ * 2) / (elapsed * 1000000) # 2 ops per iteration\n return {\n\ + \ \"iterations\": iterations,\n \"time_seconds\": f\"\ + {elapsed:.4f}\",\n \"mflops\": f\"{mflops:.2f}\"\n }\n\n\ + def memory_bandwidth_benchmark():\n \"\"\"Memory sequential read/write\ + \ benchmark\"\"\"\n size_mb = 100\n size_bytes = size_mb * 1024\ + \ * 1024\n data = bytearray(size_bytes)\n\n # Write test\n \ + \ start = time.time()\n for i in range(0, size_bytes, 8):\n \ + \ data[i:i+8] = b'\\x00\\x01\\x02\\x03\\x04\\x05\\x06\\x07'\n\ + \ write_time = time.time() - start\n write_bw = (size_mb / write_time)\ + \ if write_time > 0 else 0\n\n # Read test\n start = time.time()\n\ + \ checksum = 0\n for i in range(0, size_bytes, 8):\n \ + \ checksum += data[i]\n read_time = time.time() - start\n read_bw\ + \ = (size_mb / read_time) if read_time > 0 else 0\n\n return {\n\ + \ \"test_size_mb\": size_mb,\n \"write_bandwidth_mbps\"\ + : f\"{write_bw:.2f}\",\n \"read_bandwidth_mbps\": f\"{read_bw:.2f}\"\ + ,\n \"write_time_sec\": f\"{write_time:.4f}\",\n \"\ + read_time_sec\": f\"{read_time:.4f}\"\n }\n\ndef disk_io_benchmark():\n\ + \ \"\"\"Disk I/O benchmark using /tmp\"\"\"\n test_file = \"\ + /tmp/io_bench_test.dat\"\n size_mb = 50\n block_size = 1024\ + \ * 1024 # 1MB blocks\n\n try:\n # Write test\n \ + \ start = time.time()\n with open(test_file, 'wb') as f:\n\ + \ for _ in range(size_mb):\n f.write(os.urandom(block_size))\n\ + \ f.flush()\n os.fsync(f.fileno())\n \ + \ write_time = time.time() - start\n write_bw = (size_mb /\ + \ write_time) if write_time > 0 else 0\n\n # Read test\n \ + \ start = time.time()\n with open(test_file, 'rb') as\ + \ f:\n while f.read(block_size):\n pass\n\ + \ read_time = time.time() - start\n read_bw = (size_mb\ + \ / read_time) if read_time > 0 else 0\n\n # Cleanup\n \ + \ os.remove(test_file)\n\n return {\n \"test_size_mb\"\ + : size_mb,\n \"write_bandwidth_mbps\": f\"{write_bw:.2f}\"\ + ,\n \"read_bandwidth_mbps\": f\"{read_bw:.2f}\",\n \ + \ \"write_time_sec\": f\"{write_time:.4f}\",\n \"\ + read_time_sec\": f\"{read_time:.4f}\"\n }\n except Exception\ + \ as e:\n return {\"error\": str(e)}\n\ndef context_switch_benchmark():\n\ + \ \"\"\"Measure context switch overhead using subprocess spawning\"\ + \"\"\n iterations = 100\n start = time.time()\n for _ in\ + \ range(iterations):\n subprocess.run(['true'], capture_output=True)\n\ + \ elapsed = time.time() - start\n avg_per_switch = (elapsed\ + \ / iterations) * 1000 # milliseconds\n\n return {\n \"\ + iterations\": iterations,\n \"total_time_sec\": f\"{elapsed:.4f}\"\ + ,\n \"avg_per_spawn_ms\": f\"{avg_per_switch:.4f}\"\n }\n\ + \nprint(\"Collecting static system information...\")\n\nsnapshot =\ + \ {\n \"timestamp\": datetime.now().isoformat(),\n \"system\"\ + : {\n \"cpu_model\": run_cmd(\"grep 'model name' /proc/cpuinfo\ + \ | head -1 | cut -d: -f2 | xargs\"),\n \"cpu_count\": run_cmd(\"\ + grep -c ^processor /proc/cpuinfo\"),\n \"cpu_cores_per_socket\"\ + : run_cmd(\"lscpu | grep 'Core(s) per socket' | awk '{print $NF}'\ + \ 2>/dev/null || echo 'N/A'\"),\n \"cpu_threads_per_core\"\ + : run_cmd(\"lscpu | grep 'Thread(s) per core' | awk '{print $NF}'\ + \ 2>/dev/null || echo 'N/A'\"),\n \"cpu_architecture\": run_cmd(\"\ + uname -m\"),\n \"cpu_current_freq_mhz\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_cur_freq\ + \ 2>/dev/null | awk '{print $1/1000}' || echo 'N/A'\"),\n \"\ + cpu_max_freq_mhz\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_max_freq\ + \ 2>/dev/null | awk '{print $1/1000}' || echo 'N/A'\"),\n \"\ + cpu_min_freq_mhz\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_min_freq\ + \ 2>/dev/null | awk '{print $1/1000}' || echo 'N/A'\"),\n \"\ + cache_l1d\": run_cmd(\"lscpu | grep 'L1d cache' | awk '{print $NF}'\ + \ 2>/dev/null || echo 'N/A'\"),\n \"cache_l1i\": run_cmd(\"\ + lscpu | grep 'L1i cache' | awk '{print $NF}' 2>/dev/null || echo 'N/A'\"\ + ),\n \"cache_l2\": run_cmd(\"lscpu | grep 'L2 cache' | awk\ + \ '{print $NF}' 2>/dev/null || echo 'N/A'\"),\n \"cache_l3\"\ + : run_cmd(\"lscpu | grep 'L3 cache' | awk '{print $NF}' 2>/dev/null\ + \ || echo 'N/A'\"),\n \"memory_total_gb\": run_cmd(\"grep MemTotal\ + \ /proc/meminfo | awk '{printf \\\"%.2f\\\", $2/1024/1024}'\"),\n\ + \ \"memory_available_gb\": run_cmd(\"grep MemAvailable /proc/meminfo\ + \ | awk '{printf \\\"%.2f\\\", $2/1024/1024}'\"),\n \"memory_free_gb\"\ + : run_cmd(\"grep MemFree /proc/meminfo | awk '{printf \\\"%.2f\\\"\ + , $2/1024/1024}'\"),\n \"numa_node_count\": run_cmd(\"ls -d\ + \ /sys/devices/system/node/node* 2>/dev/null | wc -l\"),\n \ + \ \"kernel_version\": run_cmd(\"uname -r\"),\n \"uptime_seconds\"\ + : run_cmd(\"cat /proc/uptime | awk '{print $1}'\")\n },\n \"\ + gpu\": {\n \"gpu_count\": run_cmd(\"nvidia-smi --query-gpu=name\ + \ --format=csv,noheader 2>/dev/null | wc -l\"),\n \"gpu_models\"\ + : run_cmd(\"nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null\"\ + ),\n \"gpu_memory_per_device_mb\": run_cmd(\"nvidia-smi --query-gpu=memory.total\ + \ --format=csv,noheader,nounits 2>/dev/null\"),\n \"gpu_driver_version\"\ + : run_cmd(\"nvidia-smi --query-gpu=driver_version --format=csv,noheader\ + \ 2>/dev/null | head -1\"),\n \"cuda_version\": run_cmd(\"\ + nvidia-smi | grep 'CUDA Version' | awk '{print $9}' 2>/dev/null ||\ + \ echo 'N/A'\")\n },\n \"network\": {\n \"interface_names\"\ + : run_cmd(\"ls /sys/class/net 2>/dev/null | grep -v lo | tr '\\\\\ + n' ', ' | sed 's/,$//'\"),\n \"rdma_devices\": run_cmd(\"ls\ + \ /sys/class/infiniband 2>/dev/null | tr '\\\\n' ', ' | sed 's/,$//'\ + \ || echo 'None'\"),\n \"rdma_device_count\": run_cmd(\"ls\ + \ /sys/class/infiniband 2>/dev/null | wc -l\"),\n \"primary_mtu\"\ + : run_cmd(\"cat /sys/class/net/eth0/mtu 2>/dev/null || echo 'N/A'\"\ + ),\n \"tcp_rmem\": run_cmd(\"cat /proc/sys/net/ipv4/tcp_rmem\ + \ 2>/dev/null || echo 'N/A'\"),\n \"tcp_wmem\": run_cmd(\"\ + cat /proc/sys/net/ipv4/tcp_wmem 2>/dev/null || echo 'N/A'\"),\n \ + \ \"tcp_congestion_control\": run_cmd(\"cat /proc/sys/net/ipv4/tcp_congestion_control\ + \ 2>/dev/null || echo 'N/A'\")\n },\n \"storage\": {\n \ + \ \"block_devices\": run_cmd(\"lsblk -nd -o NAME,SIZE 2>/dev/null\ + \ | grep -v loop\"),\n \"nvme_count\": run_cmd(\"ls /dev/nvme*\ + \ 2>/dev/null | grep -c nvme0n || echo '0'\"),\n \"tmpfs_size\"\ + : run_cmd(\"df -h /tmp 2>/dev/null | tail -1 | awk '{print $2}'\"\ + ),\n \"tmpfs_avail\": run_cmd(\"df -h /tmp 2>/dev/null | tail\ + \ -1 | awk '{print $4}'\"),\n \"io_scheduler\": run_cmd(\"\ + cat /sys/block/sda/queue/scheduler 2>/dev/null | grep -o '\\\\[.*\\\ + \\]' | tr -d '[]' || echo 'N/A'\")\n },\n \"performance_config\"\ + : {\n \"cpu_governor\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_governor\ + \ 2>/dev/null || echo 'N/A'\"),\n \"numa_balancing\": run_cmd(\"\ + cat /proc/sys/kernel/numa_balancing 2>/dev/null || echo 'N/A'\"),\n\ + \ \"transparent_hugepages\": run_cmd(\"cat /sys/kernel/mm/transparent_hugepage/enabled\ + \ 2>/dev/null | grep -o '\\\\[.*\\\\]' | tr -d '[]' || echo 'N/A'\"\ + ),\n \"swappiness\": run_cmd(\"cat /proc/sys/vm/swappiness\ + \ 2>/dev/null || echo 'N/A'\"),\n \"dirty_ratio\": run_cmd(\"\ + cat /proc/sys/vm/dirty_ratio 2>/dev/null || echo 'N/A'\"),\n \ + \ \"dirty_background_ratio\": run_cmd(\"cat /proc/sys/vm/dirty_background_ratio\ + \ 2>/dev/null || echo 'N/A'\"),\n \"max_map_count\": run_cmd(\"\ + cat /proc/sys/vm/max_map_count 2>/dev/null || echo 'N/A'\"),\n \ + \ \"file_max\": run_cmd(\"cat /proc/sys/fs/file-max 2>/dev/null\ + \ || echo 'N/A'\")\n },\n \"process_limits\": {\n \"\ + max_user_processes\": run_cmd(\"ulimit -u 2>/dev/null || echo 'N/A'\"\ + ),\n \"max_open_files\": run_cmd(\"ulimit -n 2>/dev/null ||\ + \ echo 'N/A'\"),\n \"max_stack_size_kb\": run_cmd(\"ulimit\ + \ -s 2>/dev/null || echo 'N/A'\"),\n \"max_memory_size_kb\"\ + : run_cmd(\"ulimit -m 2>/dev/null || echo 'N/A'\"),\n \"cgroup_cpu_quota\"\ + : run_cmd(\"cat /sys/fs/cgroup/cpu/cpu.cfs_quota_us 2>/dev/null ||\ + \ echo 'N/A'\"),\n \"cgroup_cpu_period\": run_cmd(\"cat /sys/fs/cgroup/cpu/cpu.cfs_period_us\ + \ 2>/dev/null || echo 'N/A'\"),\n \"cgroup_memory_limit_bytes\"\ + : run_cmd(\"cat /sys/fs/cgroup/memory/memory.limit_in_bytes 2>/dev/null\ + \ || echo 'N/A'\")\n }\n}\n\n# Dynamic Benchmarks\nprint(\"Running\ + \ CPU compute benchmark...\")\nsnapshot[\"benchmarks\"] = {\n \"\ + cpu_compute\": cpu_benchmark()\n}\n\nprint(\"Running memory bandwidth\ + \ benchmark...\")\nsnapshot[\"benchmarks\"][\"memory_bandwidth\"]\ + \ = memory_bandwidth_benchmark()\n\nprint(\"Running disk I/O benchmark...\"\ + )\nsnapshot[\"benchmarks\"][\"disk_io\"] = disk_io_benchmark()\n\n\ + print(\"Running context switch benchmark...\")\nsnapshot[\"benchmarks\"\ + ][\"context_switch\"] = context_switch_benchmark()\n\n# Write JSON\n\ + with open(\"/tmp/result/discovery.json\", \"w\") as f:\n json.dump(snapshot,\ + \ f, indent=2)\n\nprint(\"Snapshot generation complete!\")\nEOF_SNAPSHOT\n\ + \n# Execute snapshot generation\nmkdir -p /tmp/result\npython3 /tmp/generate_snapshot.py\n\ + \n# Clean up intermediate scripts\nrm -f /tmp/generate_snapshot.py\n\ + \necho \"User system snapshot complete. Results in discovery.json\"\ + \n\nif [ -f /tmp/result/discovery.json ]; then mv /tmp/result/discovery.json\ + \ /tmp/result/discovery_user_snapshot.json; fi" + volumeMounts: + - name: results + mountPath: /tmp/result + subPath: coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821/task-0 + dependsOn: + - name: mkdir + status: Complete + - name: task-1 + replicas: 1 + template: + spec: + parallelism: 1 + completions: 1 + backoffLimit: 0 + template: + metadata: + labels: + app: task-1 + coldpress/gid: coldpress-ddp-training + spec: + restartPolicy: Never + containers: + - name: training + image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime + workingDir: /workspace + command: + - bash + - -c + args: + - '# Build torch.distributed.run command from env vars + + CMD="python -m torch.distributed.run" + + + # Add distributed params if env vars are set + + [ -n "$NPROC_PER_NODE" ] && CMD="$CMD --nproc_per_node=$NPROC_PER_NODE" + + [ -n "$NNODES" ] && CMD="$CMD --nnodes=$NNODES" + + [ -n "$RANK" ] && CMD="$CMD --node_rank=$RANK" + + [ -n "$MASTER_ADDR" ] && CMD="$CMD --master_addr=$MASTER_ADDR" + + [ -n "$MASTER_PORT" ] && CMD="$CMD --master_port=$MASTER_PORT" + + + # Training script and args + + CMD="$CMD train.py" + + CMD="$CMD --dataset=mnist" + + CMD="$CMD --train-test-split=0.8" + + CMD="$CMD --epochs=50" + + CMD="$CMD --batch-size=128" + + CMD="$CMD --hidden-size=4096" + + CMD="$CMD --lr=0.01" + + CMD="$CMD --output-dir=/results/checkpoints" + + + echo "Running: $CMD" + + exec $CMD + + ' + resources: + requests: + nvidia.com/gpu: '2' + memory: 16Gi + cpu: '8' + limits: + nvidia.com/gpu: '2' + memory: 16Gi + env: + - name: NCCL_DEBUG + value: INFO + - name: NPROC_PER_NODE + value: '2' + - name: NNODES + value: '2' + - name: MASTER_ADDR + value: coldpress-ddp-training-task-0-0-0.coldpress-ddp-training.coldpress-project.svc.cluster.local + - name: MASTER_PORT + value: '29500' + - name: RANK + value: '1' + - name: WORLD_SIZE + value: '2' + volumeMounts: + - name: results + mountPath: /results + subPath: coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821/task-1 + - name: dshm + mountPath: /dev/shm + - name: training-script + mountPath: /workspace/train.py + subPath: train.py + tolerations: *id001 + volumes: + - name: results + persistentVolumeClaim: + claimName: coldpress-project-storage + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 16Gi + - name: training-script + configMap: + name: coldpress-ddp-training-files + initContainers: + - name: discovery + image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime + command: + - /bin/bash + - -c + args: + - "echo \"User System Snapshot - Capturing performance baseline with\ + \ dynamic benchmarks...\"\n\n# Create comprehensive Python script\ + \ in /tmp\ncat << 'EOF_SNAPSHOT' > /tmp/generate_snapshot.py\nimport\ + \ subprocess\nimport json\nimport time\nimport os\nfrom datetime import\ + \ datetime\n\ndef run_cmd(command):\n try:\n out = subprocess.run(command,\ + \ shell=True, capture_output=True, text=True, timeout=10)\n \ + \ return out.stdout.strip() if out.returncode == 0 else \"Not available\"\ + \n except Exception as e:\n return \"Not available\"\n\n\ + def cpu_benchmark():\n \"\"\"Simple CPU compute benchmark - measures\ + \ MFLOPS\"\"\"\n start = time.time()\n result = 0.0\n iterations\ + \ = 10000000\n for i in range(iterations):\n result += i\ + \ * 1.5 - 0.7\n elapsed = time.time() - start\n mflops = (iterations\ + \ * 2) / (elapsed * 1000000) # 2 ops per iteration\n return {\n\ + \ \"iterations\": iterations,\n \"time_seconds\": f\"\ + {elapsed:.4f}\",\n \"mflops\": f\"{mflops:.2f}\"\n }\n\n\ + def memory_bandwidth_benchmark():\n \"\"\"Memory sequential read/write\ + \ benchmark\"\"\"\n size_mb = 100\n size_bytes = size_mb * 1024\ + \ * 1024\n data = bytearray(size_bytes)\n\n # Write test\n \ + \ start = time.time()\n for i in range(0, size_bytes, 8):\n \ + \ data[i:i+8] = b'\\x00\\x01\\x02\\x03\\x04\\x05\\x06\\x07'\n\ + \ write_time = time.time() - start\n write_bw = (size_mb / write_time)\ + \ if write_time > 0 else 0\n\n # Read test\n start = time.time()\n\ + \ checksum = 0\n for i in range(0, size_bytes, 8):\n \ + \ checksum += data[i]\n read_time = time.time() - start\n read_bw\ + \ = (size_mb / read_time) if read_time > 0 else 0\n\n return {\n\ + \ \"test_size_mb\": size_mb,\n \"write_bandwidth_mbps\"\ + : f\"{write_bw:.2f}\",\n \"read_bandwidth_mbps\": f\"{read_bw:.2f}\"\ + ,\n \"write_time_sec\": f\"{write_time:.4f}\",\n \"\ + read_time_sec\": f\"{read_time:.4f}\"\n }\n\ndef disk_io_benchmark():\n\ + \ \"\"\"Disk I/O benchmark using /tmp\"\"\"\n test_file = \"\ + /tmp/io_bench_test.dat\"\n size_mb = 50\n block_size = 1024\ + \ * 1024 # 1MB blocks\n\n try:\n # Write test\n \ + \ start = time.time()\n with open(test_file, 'wb') as f:\n\ + \ for _ in range(size_mb):\n f.write(os.urandom(block_size))\n\ + \ f.flush()\n os.fsync(f.fileno())\n \ + \ write_time = time.time() - start\n write_bw = (size_mb /\ + \ write_time) if write_time > 0 else 0\n\n # Read test\n \ + \ start = time.time()\n with open(test_file, 'rb') as\ + \ f:\n while f.read(block_size):\n pass\n\ + \ read_time = time.time() - start\n read_bw = (size_mb\ + \ / read_time) if read_time > 0 else 0\n\n # Cleanup\n \ + \ os.remove(test_file)\n\n return {\n \"test_size_mb\"\ + : size_mb,\n \"write_bandwidth_mbps\": f\"{write_bw:.2f}\"\ + ,\n \"read_bandwidth_mbps\": f\"{read_bw:.2f}\",\n \ + \ \"write_time_sec\": f\"{write_time:.4f}\",\n \"\ + read_time_sec\": f\"{read_time:.4f}\"\n }\n except Exception\ + \ as e:\n return {\"error\": str(e)}\n\ndef context_switch_benchmark():\n\ + \ \"\"\"Measure context switch overhead using subprocess spawning\"\ + \"\"\n iterations = 100\n start = time.time()\n for _ in\ + \ range(iterations):\n subprocess.run(['true'], capture_output=True)\n\ + \ elapsed = time.time() - start\n avg_per_switch = (elapsed\ + \ / iterations) * 1000 # milliseconds\n\n return {\n \"\ + iterations\": iterations,\n \"total_time_sec\": f\"{elapsed:.4f}\"\ + ,\n \"avg_per_spawn_ms\": f\"{avg_per_switch:.4f}\"\n }\n\ + \nprint(\"Collecting static system information...\")\n\nsnapshot =\ + \ {\n \"timestamp\": datetime.now().isoformat(),\n \"system\"\ + : {\n \"cpu_model\": run_cmd(\"grep 'model name' /proc/cpuinfo\ + \ | head -1 | cut -d: -f2 | xargs\"),\n \"cpu_count\": run_cmd(\"\ + grep -c ^processor /proc/cpuinfo\"),\n \"cpu_cores_per_socket\"\ + : run_cmd(\"lscpu | grep 'Core(s) per socket' | awk '{print $NF}'\ + \ 2>/dev/null || echo 'N/A'\"),\n \"cpu_threads_per_core\"\ + : run_cmd(\"lscpu | grep 'Thread(s) per core' | awk '{print $NF}'\ + \ 2>/dev/null || echo 'N/A'\"),\n \"cpu_architecture\": run_cmd(\"\ + uname -m\"),\n \"cpu_current_freq_mhz\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_cur_freq\ + \ 2>/dev/null | awk '{print $1/1000}' || echo 'N/A'\"),\n \"\ + cpu_max_freq_mhz\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_max_freq\ + \ 2>/dev/null | awk '{print $1/1000}' || echo 'N/A'\"),\n \"\ + cpu_min_freq_mhz\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_min_freq\ + \ 2>/dev/null | awk '{print $1/1000}' || echo 'N/A'\"),\n \"\ + cache_l1d\": run_cmd(\"lscpu | grep 'L1d cache' | awk '{print $NF}'\ + \ 2>/dev/null || echo 'N/A'\"),\n \"cache_l1i\": run_cmd(\"\ + lscpu | grep 'L1i cache' | awk '{print $NF}' 2>/dev/null || echo 'N/A'\"\ + ),\n \"cache_l2\": run_cmd(\"lscpu | grep 'L2 cache' | awk\ + \ '{print $NF}' 2>/dev/null || echo 'N/A'\"),\n \"cache_l3\"\ + : run_cmd(\"lscpu | grep 'L3 cache' | awk '{print $NF}' 2>/dev/null\ + \ || echo 'N/A'\"),\n \"memory_total_gb\": run_cmd(\"grep MemTotal\ + \ /proc/meminfo | awk '{printf \\\"%.2f\\\", $2/1024/1024}'\"),\n\ + \ \"memory_available_gb\": run_cmd(\"grep MemAvailable /proc/meminfo\ + \ | awk '{printf \\\"%.2f\\\", $2/1024/1024}'\"),\n \"memory_free_gb\"\ + : run_cmd(\"grep MemFree /proc/meminfo | awk '{printf \\\"%.2f\\\"\ + , $2/1024/1024}'\"),\n \"numa_node_count\": run_cmd(\"ls -d\ + \ /sys/devices/system/node/node* 2>/dev/null | wc -l\"),\n \ + \ \"kernel_version\": run_cmd(\"uname -r\"),\n \"uptime_seconds\"\ + : run_cmd(\"cat /proc/uptime | awk '{print $1}'\")\n },\n \"\ + gpu\": {\n \"gpu_count\": run_cmd(\"nvidia-smi --query-gpu=name\ + \ --format=csv,noheader 2>/dev/null | wc -l\"),\n \"gpu_models\"\ + : run_cmd(\"nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null\"\ + ),\n \"gpu_memory_per_device_mb\": run_cmd(\"nvidia-smi --query-gpu=memory.total\ + \ --format=csv,noheader,nounits 2>/dev/null\"),\n \"gpu_driver_version\"\ + : run_cmd(\"nvidia-smi --query-gpu=driver_version --format=csv,noheader\ + \ 2>/dev/null | head -1\"),\n \"cuda_version\": run_cmd(\"\ + nvidia-smi | grep 'CUDA Version' | awk '{print $9}' 2>/dev/null ||\ + \ echo 'N/A'\")\n },\n \"network\": {\n \"interface_names\"\ + : run_cmd(\"ls /sys/class/net 2>/dev/null | grep -v lo | tr '\\\\\ + n' ', ' | sed 's/,$//'\"),\n \"rdma_devices\": run_cmd(\"ls\ + \ /sys/class/infiniband 2>/dev/null | tr '\\\\n' ', ' | sed 's/,$//'\ + \ || echo 'None'\"),\n \"rdma_device_count\": run_cmd(\"ls\ + \ /sys/class/infiniband 2>/dev/null | wc -l\"),\n \"primary_mtu\"\ + : run_cmd(\"cat /sys/class/net/eth0/mtu 2>/dev/null || echo 'N/A'\"\ + ),\n \"tcp_rmem\": run_cmd(\"cat /proc/sys/net/ipv4/tcp_rmem\ + \ 2>/dev/null || echo 'N/A'\"),\n \"tcp_wmem\": run_cmd(\"\ + cat /proc/sys/net/ipv4/tcp_wmem 2>/dev/null || echo 'N/A'\"),\n \ + \ \"tcp_congestion_control\": run_cmd(\"cat /proc/sys/net/ipv4/tcp_congestion_control\ + \ 2>/dev/null || echo 'N/A'\")\n },\n \"storage\": {\n \ + \ \"block_devices\": run_cmd(\"lsblk -nd -o NAME,SIZE 2>/dev/null\ + \ | grep -v loop\"),\n \"nvme_count\": run_cmd(\"ls /dev/nvme*\ + \ 2>/dev/null | grep -c nvme0n || echo '0'\"),\n \"tmpfs_size\"\ + : run_cmd(\"df -h /tmp 2>/dev/null | tail -1 | awk '{print $2}'\"\ + ),\n \"tmpfs_avail\": run_cmd(\"df -h /tmp 2>/dev/null | tail\ + \ -1 | awk '{print $4}'\"),\n \"io_scheduler\": run_cmd(\"\ + cat /sys/block/sda/queue/scheduler 2>/dev/null | grep -o '\\\\[.*\\\ + \\]' | tr -d '[]' || echo 'N/A'\")\n },\n \"performance_config\"\ + : {\n \"cpu_governor\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_governor\ + \ 2>/dev/null || echo 'N/A'\"),\n \"numa_balancing\": run_cmd(\"\ + cat /proc/sys/kernel/numa_balancing 2>/dev/null || echo 'N/A'\"),\n\ + \ \"transparent_hugepages\": run_cmd(\"cat /sys/kernel/mm/transparent_hugepage/enabled\ + \ 2>/dev/null | grep -o '\\\\[.*\\\\]' | tr -d '[]' || echo 'N/A'\"\ + ),\n \"swappiness\": run_cmd(\"cat /proc/sys/vm/swappiness\ + \ 2>/dev/null || echo 'N/A'\"),\n \"dirty_ratio\": run_cmd(\"\ + cat /proc/sys/vm/dirty_ratio 2>/dev/null || echo 'N/A'\"),\n \ + \ \"dirty_background_ratio\": run_cmd(\"cat /proc/sys/vm/dirty_background_ratio\ + \ 2>/dev/null || echo 'N/A'\"),\n \"max_map_count\": run_cmd(\"\ + cat /proc/sys/vm/max_map_count 2>/dev/null || echo 'N/A'\"),\n \ + \ \"file_max\": run_cmd(\"cat /proc/sys/fs/file-max 2>/dev/null\ + \ || echo 'N/A'\")\n },\n \"process_limits\": {\n \"\ + max_user_processes\": run_cmd(\"ulimit -u 2>/dev/null || echo 'N/A'\"\ + ),\n \"max_open_files\": run_cmd(\"ulimit -n 2>/dev/null ||\ + \ echo 'N/A'\"),\n \"max_stack_size_kb\": run_cmd(\"ulimit\ + \ -s 2>/dev/null || echo 'N/A'\"),\n \"max_memory_size_kb\"\ + : run_cmd(\"ulimit -m 2>/dev/null || echo 'N/A'\"),\n \"cgroup_cpu_quota\"\ + : run_cmd(\"cat /sys/fs/cgroup/cpu/cpu.cfs_quota_us 2>/dev/null ||\ + \ echo 'N/A'\"),\n \"cgroup_cpu_period\": run_cmd(\"cat /sys/fs/cgroup/cpu/cpu.cfs_period_us\ + \ 2>/dev/null || echo 'N/A'\"),\n \"cgroup_memory_limit_bytes\"\ + : run_cmd(\"cat /sys/fs/cgroup/memory/memory.limit_in_bytes 2>/dev/null\ + \ || echo 'N/A'\")\n }\n}\n\n# Dynamic Benchmarks\nprint(\"Running\ + \ CPU compute benchmark...\")\nsnapshot[\"benchmarks\"] = {\n \"\ + cpu_compute\": cpu_benchmark()\n}\n\nprint(\"Running memory bandwidth\ + \ benchmark...\")\nsnapshot[\"benchmarks\"][\"memory_bandwidth\"]\ + \ = memory_bandwidth_benchmark()\n\nprint(\"Running disk I/O benchmark...\"\ + )\nsnapshot[\"benchmarks\"][\"disk_io\"] = disk_io_benchmark()\n\n\ + print(\"Running context switch benchmark...\")\nsnapshot[\"benchmarks\"\ + ][\"context_switch\"] = context_switch_benchmark()\n\n# Write JSON\n\ + with open(\"/tmp/result/discovery.json\", \"w\") as f:\n json.dump(snapshot,\ + \ f, indent=2)\n\nprint(\"Snapshot generation complete!\")\nEOF_SNAPSHOT\n\ + \n# Execute snapshot generation\nmkdir -p /tmp/result\npython3 /tmp/generate_snapshot.py\n\ + \n# Clean up intermediate scripts\nrm -f /tmp/generate_snapshot.py\n\ + \necho \"User system snapshot complete. Results in discovery.json\"\ + \n\nif [ -f /tmp/result/discovery.json ]; then mv /tmp/result/discovery.json\ + \ /tmp/result/discovery_user_snapshot.json; fi" + volumeMounts: + - name: results + mountPath: /tmp/result + subPath: coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821/task-1 + dependsOn: + - name: mkdir + status: Complete + successPolicy: + operator: All + targetReplicatedJobs: + - task-0 + - task-1 diff --git a/examples/pytorch_ddp_training/comp/pytorchjob.yaml b/examples/pytorch_ddp_training/comp/pytorchjob.yaml new file mode 100644 index 0000000..d8c5b22 --- /dev/null +++ b/examples/pytorch_ddp_training/comp/pytorchjob.yaml @@ -0,0 +1,131 @@ +apiVersion: kubeflow.org/v1 +kind: PyTorchJob +metadata: + name: coldpress-ddp-training + namespace: coldpress-project + labels: + app.kubernetes.io/managed-by: coldpress + app.kubernetes.io/version: 0.2.0 + coldpress.io/job-id: coldpress-ddp-training + kueue.x-k8s.io/queue-name: coldpress-local-queue-coldpress-project + annotations: + coldpress.io/base-dir: coldpress-project/coldpress_results/ddp-training-050895d3-20260416_000443 + coldpress.io/storage-pvc: coldpress-project-storage + kueue.x-k8s.io/wait-for-pods-ready-timeout: 20m +spec: + suspend: true + pytorchReplicaSpecs: + Master: + replicas: 1 + restartPolicy: OnFailure + template: &id001 + metadata: + annotations: + sidecar.istio.io/inject: 'false' + labels: + app.kubernetes.io/managed-by: coldpress + app.kubernetes.io/version: 0.2.0 + coldpress.io/job-id: coldpress-ddp-training + spec: + restartPolicy: OnFailure + containers: + - name: pytorch + image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime + workingDir: /workspace + command: + - bash + - -c + args: + - '# Build torch.distributed.run command from env vars + + CMD="python -m torch.distributed.run" + + + # Add distributed params if env vars are set + + [ -n "$NPROC_PER_NODE" ] && CMD="$CMD --nproc_per_node=$NPROC_PER_NODE" + + [ -n "$NNODES" ] && CMD="$CMD --nnodes=$NNODES" + + [ -n "$RANK" ] && CMD="$CMD --node_rank=$RANK" + + [ -n "$MASTER_ADDR" ] && CMD="$CMD --master_addr=$MASTER_ADDR" + + [ -n "$MASTER_PORT" ] && CMD="$CMD --master_port=$MASTER_PORT" + + + # Training script and args + + CMD="$CMD train.py" + + CMD="$CMD --dataset=mnist" + + CMD="$CMD --train-test-split=0.8" + + CMD="$CMD --epochs=50" + + CMD="$CMD --batch-size=128" + + CMD="$CMD --hidden-size=4096" + + CMD="$CMD --lr=0.01" + + CMD="$CMD --output-dir=/results/checkpoints" + + + echo "Running: $CMD" + + exec $CMD + + ' + - --nnodes=2 + resources: + requests: + nvidia.com/gpu: '2' + memory: 16Gi + cpu: '8' + limits: + nvidia.com/gpu: '2' + memory: 16Gi + env: + - name: NCCL_DEBUG + value: INFO + - name: NPROC_PER_NODE + value: '2' + volumeMounts: + - name: results + mountPath: /results + - name: dshm + mountPath: /dev/shm + - name: training-script + mountPath: /workspace/train.py + subPath: train.py + tolerations: + - operator: Exists + volumes: + - name: results + persistentVolumeClaim: + claimName: coldpress-project-storage + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 16Gi + - name: training-script + configMap: + name: coldpress-ddp-training-files + initContainers: + - name: mkdir + image: registry.access.redhat.com/ubi9/ubi-minimal:latest + command: + - sh + - -c + args: + - mkdir -p /results/coldpress-project/coldpress_results/ddp-training-050895d3-20260416_000443 + && echo Created directory /results/coldpress-project/coldpress_results/ddp-training-050895d3-20260416_000443 + volumeMounts: + - name: results + mountPath: /results + Worker: + replicas: 1 + restartPolicy: OnFailure + template: *id001 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..ddfbd6c --- /dev/null +++ b/examples/pytorch_ddp_training/intent_jobset.yaml @@ -0,0 +1,21 @@ +project: coldpress-project +output: ddp-training-job +target: jobset + +files: + - train.py + +discovery: + template: user_snapshot + tasks: all + +tasks: + - name: ddp-training + replicas: 2 + env: + # Don't override NPROC_PER_NODE - use value from job-spec (2 GPUs per node) + NNODES: "${REPLICAS}" + MASTER_ADDR: "${REPLICA_ddp-training_0}" + MASTER_PORT: "29500" + RANK: "${INDEX}" + WORLD_SIZE: "${REPLICAS}" diff --git a/examples/pytorch_ddp_training/intent_kubeflow.yaml b/examples/pytorch_ddp_training/intent_kubeflow.yaml new file mode 100644 index 0000000..a4341ab --- /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: + nnodes: "${REPLICAS}" + # Kubeflow PyTorchJob auto-injects: + # - MASTER_ADDR + # - MASTER_PORT + # - RANK (via RANK env var) + # So we don't need to specify them here diff --git a/examples/pytorch_ddp_training/job-spec.yaml b/examples/pytorch_ddp_training/job-spec.yaml index 0f198f9..7d3229c 100644 --- a/examples/pytorch_ddp_training/job-spec.yaml +++ b/examples/pytorch_ddp_training/job-spec.yaml @@ -1,41 +1,75 @@ -# Generated by: Claude Sonnet 4.5 -name: ddp-training +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: ["bash", "-c"] + args: + - | + # Build torch.distributed.run command from env vars + CMD="python -m torch.distributed.run" -tolerate_all: true + # Add distributed params if env vars are set + [ -n "$NPROC_PER_NODE" ] && CMD="$CMD --nproc_per_node=$NPROC_PER_NODE" + [ -n "$NNODES" ] && CMD="$CMD --nnodes=$NNODES" + [ -n "$RANK" ] && CMD="$CMD --node_rank=$RANK" + [ -n "$MASTER_ADDR" ] && CMD="$CMD --master_addr=$MASTER_ADDR" + [ -n "$MASTER_PORT" ] && CMD="$CMD --master_port=$MASTER_PORT" -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" + # Training script and args + CMD="$CMD train.py" + CMD="$CMD --dataset=mnist" + CMD="$CMD --train-test-split=0.8" + CMD="$CMD --epochs=50" + CMD="$CMD --batch-size=128" + CMD="$CMD --hidden-size=4096" + CMD="$CMD --lr=0.01" + CMD="$CMD --output-dir=/results/checkpoints" -volumes: - - name: results - mount: /results - - name: dshm - type: emptyDir - medium: Memory - sizeLimit: 16Gi - mount: /dev/shm + echo "Running: $CMD" + exec $CMD + resources: + requests: + nvidia.com/gpu: "2" + memory: "16Gi" + cpu: "8" + limits: + nvidia.com/gpu: "2" + memory: "16Gi" + env: + - name: NCCL_DEBUG + value: "INFO" + - name: NPROC_PER_NODE + value: "2" + 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/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..e59aa85 --- /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 + env: + GUIDELLM_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..74a687e 100644 --- a/examples/vllm_guidellm_benchmark/job-spec.yaml +++ b/examples/vllm_guidellm_benchmark/job-spec.yaml @@ -1,59 +1,117 @@ -# 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: ["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 -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: ["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://localhost: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/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..ba6da1e 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, or both (comma-separated) +targets: jobset,kubeflow + # References to cluster-wide resources cluster_queue: coldpress-cluster-queue storage_class: nfs-csi 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_labels.py b/tests/test_labels.py index 08b295c..91f3983 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(): @@ -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" 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(): 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]}...") From b07c53db217e46c2538363504b50832dd1024100 Mon Sep 17 00:00:00 2001 From: asanaullah Date: Wed, 29 Apr 2026 22:04:26 -0700 Subject: [PATCH 02/14] checkpoint: pre-refactoring snapshot (v0.2.1 with kuberay support) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit captures the working state before major library-based refactoring: - JobSet, Kubeflow, KServe, and KubeRay generator support - Intent-based generation system - Custom Pydantic models and manifest builders - String-based bash script generation Planned refactoring: 1. Migrate bash script generation to Jinja2 templates (~550 line reduction) 2. Replace manual manifest dict building with kr8/k8s-client (~400 line reduction) 3. Use official kubernetes-client models (~200 line reduction) Total codebase: ~3000 lines → target: ~1500-1800 lines Co-Authored-By: Claude Sonnet 4.5 --- .github/workflows/tests.yml | 4 + .gitignore | 4 + .pre-commit-config.yaml | 50 +- cluster/ocp-test-nerc-mghpcc.yaml | 8 +- coldpress/cli.py | 281 ++---- coldpress/constants.py | 49 +- coldpress/generator.py | 216 ----- coldpress/jobset_generator.py | 246 +----- coldpress/jobset_generator_legacy.py | 803 ------------------ coldpress/kubeflow_generator.py | 739 +--------------- coldpress/kuberay_generator.py | 511 +++++++++++ coldpress/script_gen.py | 215 +++-- coldpress/utils.py | 367 ++++++++ coldpress_common/__init__.py | 4 - coldpress_common/model.py | 68 +- coldpress_setup/cli.py | 8 +- coldpress_setup/generator.py | 84 +- .../ARG_OVERRIDE_README.md | 116 +++ .../pytorch_ddp_training/intent_jobset.yaml | 26 +- .../pytorch_ddp_training/intent_kubeflow.yaml | 12 +- examples/pytorch_ddp_training/job-spec.yaml | 40 +- .../ARG_OVERRIDE_README.md | 105 +++ .../pytorch_ray_training/intent_kuberay.yaml | 20 + examples/pytorch_ray_training/job-spec.yaml | 61 ++ examples/pytorch_ray_training/train.py | 256 ++++++ .../intent_jobset.yaml | 4 +- .../vllm_guidellm_benchmark/job-spec.yaml | 60 +- projects/coldpress-project.yaml | 4 +- pyproject.toml | 2 +- tests/test_labels.py | 6 +- tests/test_script_gen_security.py | 245 ++++++ tests/test_security.py | 100 ++- 32 files changed, 2246 insertions(+), 2468 deletions(-) delete mode 100644 coldpress/generator.py delete mode 100644 coldpress/jobset_generator_legacy.py create mode 100644 coldpress/kuberay_generator.py create mode 100644 coldpress/utils.py create mode 100644 examples/pytorch_ddp_training/ARG_OVERRIDE_README.md create mode 100644 examples/pytorch_ray_training/ARG_OVERRIDE_README.md create mode 100644 examples/pytorch_ray_training/intent_kuberay.yaml create mode 100644 examples/pytorch_ray_training/job-spec.yaml create mode 100644 examples/pytorch_ray_training/train.py create mode 100644 tests/test_script_gen_security.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 7c1549e..515fec5 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -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 6e3fbfe..99db0d9 100644 --- a/.gitignore +++ b/.gitignore @@ -71,6 +71,10 @@ htmlcov/ legacy/ *.legacy +# Development features (not ready for release) +coldpress/distributed/ +examples/pytorch_coldpress_run/ + # Generated output directories output/ manifests/ 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/cluster/ocp-test-nerc-mghpcc.yaml b/cluster/ocp-test-nerc-mghpcc.yaml index b12f6d7..64d3d90 100644 --- a/cluster/ocp-test-nerc-mghpcc.yaml +++ b/cluster/ocp-test-nerc-mghpcc.yaml @@ -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 diff --git a/coldpress/cli.py b/coldpress/cli.py index f41b837..167b756 100644 --- a/coldpress/cli.py +++ b/coldpress/cli.py @@ -1,5 +1,16 @@ # Generated by: Claude Sonnet 4.5 -"""Coldpress CLI - Generate JobSet and Kubeflow manifests for AI/HPC workloads.""" +"""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 @@ -12,13 +23,20 @@ jobset_to_yaml, services_to_yaml, ) -from .kubeflow_generator import kubeflow_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 from coldpress_common import ( - validate_config, validate_intent, validate_project_config, - validate_task_specs, ) from pydantic import ValidationError @@ -29,53 +47,12 @@ @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) - - # Validate config schema - try: - validated_config = validate_config(config_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 - - if not project: - raise ValueError("Must provide --project or specify 'project' in config") - - return config_data, config_dir, project, discovery, output_name, nodes_from_config - - def _load_intent(intent_path): """Load and validate intent.yaml.""" if not os.path.exists(intent_path): @@ -124,37 +101,6 @@ def _load_vk8s_jobs(job_spec_path): return jobs, other_resources -def _load_task_specs(config_dir): - """Load and validate task specifications from job-spec.yaml (legacy format).""" - job_spec_path = os.path.join(config_dir, "job-spec.yaml") - if not os.path.exists(job_spec_path): - raise FileNotFoundError(f"Could not find job-spec.yaml in {config_dir}") - - with open(job_spec_path, "r") as f: - task_specs_data = list(yaml.safe_load_all(f)) - - if not task_specs_data: - raise ValueError("Job spec file 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 - - # 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 - ] - - -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" - - def _load_project_config(project_name): """Load project configuration file.""" project_config_file = os.path.join(PROJECT_DIR, f"{project_name}.yaml") @@ -177,136 +123,6 @@ 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, - manifest, - services, - job_spec, - base_dir, - config_dir, - task_specs, - node_assignments, - manifest_type="jobset", -): - """Write all generated files to output directory.""" - os.makedirs(output_dir, exist_ok=True) - - # Write manifest YAML based on type - if manifest_type == "jobset": - manifest_file = os.path.join(output_dir, "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_dir, "pytorchjob.yaml") - with open(manifest_file, "w") as f: - f.write(kubeflow_to_yaml(manifest)) - elif manifest_type == "tfjob": - manifest_file = os.path.join(output_dir, "tfjob.yaml") - with open(manifest_file, "w") as f: - f.write(kubeflow_to_yaml(manifest)) - elif manifest_type == "mpijob": - manifest_file = os.path.join(output_dir, "mpijob.yaml") - with open(manifest_file, "w") as f: - f.write(kubeflow_to_yaml(manifest)) - elif manifest_type == "inferenceservice": - manifest_file = os.path.join(output_dir, "inferenceservice.yaml") - with open(manifest_file, "w") as f: - f.write(kubeflow_to_yaml(manifest)) - - # Write Services YAML if any (JobSet only) - 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"], - "backend": manifest_type, - "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, - job_spec["name"], # Use original job name, not prefixed - job_spec["namespace"], - storage_pvc, - base_dir, - configmap_name, - configmap_files, - manifest_type, # Pass manifest type to scripts - ) - - return output_dir - - @cli.command() @click.option( "-i", @@ -365,6 +181,7 @@ def generate(intent): backend_name_map = { "kubeflow": "Kubeflow", "kserve": "KServe", + "kuberay": "KubeRay", "jobset": "JobSet", } backend_name = backend_name_map.get(backend, "JobSet") @@ -373,11 +190,6 @@ def generate(intent): # Generate manifests based on target if backend == "kubeflow": # Generate Kubeflow manifest (PyTorchJob) - from .kubeflow_generator import ( - generate_pytorchjob_from_intent, - kubeflow_to_yaml, - ) - manifest, base_dir = generate_pytorchjob_from_intent( vk8s_jobs, intent_config, project_config, namespace ) @@ -385,16 +197,18 @@ def generate(intent): services = [] # Kubeflow doesn't use separate services elif backend == "kserve": # Generate KServe manifest (InferenceService) - from .kubeflow_generator import ( - generate_inferenceservice_from_intent, - kubeflow_to_yaml, - ) - manifest, base_dir = generate_inferenceservice_from_intent( vk8s_jobs, intent_config, project_config, namespace ) 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: # Generate JobSet manifest manifest, services, base_dir = generate_jobset_from_intent( @@ -405,6 +219,10 @@ def generate(intent): # 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: @@ -417,6 +235,10 @@ def generate(intent): 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") @@ -486,15 +308,32 @@ def generate(intent): break if configmap_name: break + 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: + configmap_name = volume["configMap"]["name"] + break # 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 = project_config.get("storage", {}).get( - "results", f"coldpress-{namespace}-storage" - ) + # Storage is validated by Pydantic - if present, 'results' is required + storage = project_config.get("storage") + if storage: + storage_pvc = storage["results"] + else: + storage_pvc = f"coldpress-{namespace}-storage" # 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-"): @@ -542,7 +381,7 @@ def generate(intent): 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 diff --git a/coldpress/constants.py b/coldpress/constants.py index 533aa0d..eaf003e 100644 --- a/coldpress/constants.py +++ b/coldpress/constants.py @@ -2,7 +2,7 @@ """Coldpress constants - centralized configuration values.""" # Version -COLDPRESS_VERSION = "0.2.0" +COLDPRESS_VERSION = "0.2.1" # Resource naming COLDPRESS_PREFIX = "coldpress" @@ -10,6 +10,8 @@ # 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 = { @@ -17,6 +19,15 @@ "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: @@ -37,3 +48,39 @@ def get_service_name(jobset_name: str, task_id: int) -> str: 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", + }, +} diff --git a/coldpress/generator.py b/coldpress/generator.py deleted file mode 100644 index ab0da6d..0000000 --- a/coldpress/generator.py +++ /dev/null @@ -1,216 +0,0 @@ -# Generated by: Claude Sonnet 4.5 -"""Common generator utilities for Coldpress.""" - -import hashlib -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 extract_container_config(task): - """Extract container configuration from task. - - Args: - task: Task specification dict - - Returns: - dict: Container configuration with keys: image, command, args, env, - working_dir, resources, gpus - """ - 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", {}), - "ports": container_def.get("ports", []), - "readinessProbe": container_def.get("readinessProbe"), - "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", {}), - "ports": [], - "readinessProbe": None, - "gpus": task.get("gpus", 0), - } - - -def build_volumes_from_task(task, job_spec): - """Build volume list from task and job spec. - - Args: - task: Task specification dict - job_spec: Job specification dict - - Returns: - list: List of volume dicts for pod spec - """ - volumes = [] - - # Add volumes from task - task_volumes = task.get("volumes", []) - for vol in task_volumes: - vol_name = vol.get("name") - vol_type = vol.get("type", "pvc") - - if vol_type == "pvc": - # Use storage config from job_spec - storage = job_spec.get("storage", {}) - pvc_name = storage.get( - vol_name, f"coldpress-{job_spec['namespace']}-{vol_name}" - ) - volumes.append( - {"name": vol_name, "persistentVolumeClaim": {"claimName": pvc_name}} - ) - elif vol_type == "emptyDir": - volumes.append( - { - "name": vol_name, - "emptyDir": { - "medium": vol.get("medium", ""), - "sizeLimit": vol.get("sizeLimit", ""), - }, - } - ) - - return volumes - - -def build_volume_mounts_from_task(task): - """Build volume mount list from task. - - Args: - task: Task specification dict - - Returns: - list: List of volume mount dicts - """ - volume_mounts = [] - - task_volumes = task.get("volumes", []) - for vol in task_volumes: - if "mount" in vol: - volume_mounts.append({"name": vol["name"], "mountPath": vol["mount"]}) - - return volume_mounts - - -def detect_job_type(task): - """Detect job type from task configuration. - - Args: - task: Task specification dict - - Returns: - str: 'training' or 'inference' - """ - container_config = extract_container_config(task) - - # Check for readinessProbe - typical of inference services - if container_config.get("readinessProbe"): - return "inference" - - # Check image for inference patterns - image = container_config.get("image", "").lower() - if any( - pattern in image - for pattern in ["serving", "inference", "kserve", "torchserve", "triton"] - ): - return "inference" - - # Default to training - return "training" - - -def detect_framework(task): - """Detect ML framework from task configuration. - - Args: - task: Task specification dict - - Returns: - str: 'pytorch', 'tensorflow', 'mpi', or None - """ - container_config = extract_container_config(task) - image = container_config.get("image", "").lower() - command = container_config.get("command", []) - - # Check image name - if "pytorch" in image: - return "pytorch" - elif "tensorflow" in image or "tf" in image: - return "tensorflow" - - # Check command - command_str = " ".join(command) if isinstance(command, list) else str(command) - if "torch" in command_str.lower(): - return "pytorch" - elif "tensorflow" in command_str.lower(): - return "tensorflow" - elif "mpirun" in command_str.lower(): - return "mpi" - - return None diff --git a/coldpress/jobset_generator.py b/coldpress/jobset_generator.py index 55c246d..9312dcf 100644 --- a/coldpress/jobset_generator.py +++ b/coldpress/jobset_generator.py @@ -8,90 +8,24 @@ Direct transformation from vk8s to JobSet with no intermediate formats. """ -import os -import re import copy -import sys +import os import yaml from datetime import datetime, timezone from .constants import ( COLDPRESS_LABELS, MKDIR_IMAGE, + DEFAULT_DISCOVERY_DIR, get_kueue_queue_label, get_jobset_name, get_pvc_name, ) - - -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 - dangerous_chars = ["`", "$", "|", ";", "&", ">", "<", "\n", "\r"] - for char in dangerous_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) +from .utils import ( + substitute_macros, + apply_arg_overrides, + apply_env_overrides, + build_discovery_init_container, +) def build_replica_macros( @@ -127,99 +61,6 @@ def build_replica_macros( return macros -def replace_or_add_arg(args: list, key: str, value: str) -> list: - """Replace or add argument in args list.""" - new_args = [] - i = 0 - found = False - - while i < len(args): - arg = str(args[i]) - - if arg.startswith(f"--{key}="): - new_args.append(f"--{key}={value}") - found = True - i += 1 - elif arg == f"--{key}": - new_args.append(arg) - new_args.append(value) - found = True - i += 2 - else: - new_args.append(arg) - i += 1 - - if not found: - new_args.append(f"--{key}={value}") - - return new_args - - -def build_discovery_init_container( - template_path: str, task_id: int, base_dir: str, storage_volume_name: str -): - """ - 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 - storage_volume_name: Name of the storage volume to mount - - 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 = copy.deepcopy(containers[0]) - 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": storage_volume_name, - "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 generate_jobset_from_intent( vk8s_jobs: dict, intent_config, project_config, namespace: str ): @@ -242,8 +83,12 @@ def generate_jobset_from_intent( job_id = f"{intent_config.tasks[0].name}-workflow" jobset_name = get_jobset_name(job_id) - storage = project_config.get("storage", {}) - data_pvc_name = storage.get("results", get_pvc_name(namespace)) + # Storage is validated by Pydantic - if present, 'results' is required + storage = project_config.get("storage") + if storage: + data_pvc_name = storage["results"] + else: + data_pvc_name = get_pvc_name(namespace) # Generate base directory for results timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") @@ -263,7 +108,7 @@ def generate_jobset_from_intent( ) # Resolve template path - discovery_dir = os.getenv("COLDPRESS_DISCOVERY_DIR", "discovery") + 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 @@ -518,56 +363,9 @@ def build_replicated_job_from_vk8s( env_item["value"], macros, context ) - # Apply arg replacements - if task_intent.args and container.get("args"): - args = ( - container["args"] - if isinstance(container["args"], list) - else [container["args"]] - ) - - for arg_key, arg_value in task_intent.args.items(): - # Validate arg key follows Kubernetes conventions - if ( - not arg_key.replace("_", "-") - .replace(".", "-") - .replace("/", "-") - .isalnum() - ): - # Allow alphanumeric, _, -, ., / in arg keys - pass # Most arg keys are valid - - # Substitute macros with validation - context = f" for arg '--{arg_key}'" - substituted_value = substitute_macros(arg_value, macros, context) - args = replace_or_add_arg(args, arg_key, substituted_value) - - container["args"] = args - - # Apply env var additions/replacements - if task_intent.env: - # 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 + # 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 = { @@ -602,7 +400,7 @@ def build_replicated_job_from_vk8s( # Add discovery init container if specified if discovery_template_path and storage_volume_name: discovery_init = build_discovery_init_container( - discovery_template_path, task_counter, base_dir, storage_volume_name + discovery_template_path, base_dir, storage_volume_name, task_counter ) if discovery_init: pod_spec["initContainers"] = [discovery_init] @@ -666,13 +464,11 @@ def build_replicated_job_from_vk8s( def jobset_to_yaml(jobset: dict) -> str: """Convert JobSet dict to YAML string.""" - import yaml - return yaml.dump(jobset, default_flow_style=False, sort_keys=False) + 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.""" - import yaml - return yaml.dump_all(services, default_flow_style=False, sort_keys=False) + return yaml.safe_dump_all(services, default_flow_style=False, sort_keys=False) diff --git a/coldpress/jobset_generator_legacy.py b/coldpress/jobset_generator_legacy.py deleted file mode 100644 index 790bc0f..0000000 --- a/coldpress/jobset_generator_legacy.py +++ /dev/null @@ -1,803 +0,0 @@ -# Generated by: Claude Sonnet 4.5 -"""JobSet YAML generation for Coldpress jobs.""" - -import os -import sys -import yaml -import shlex -from urllib.parse import urlparse - -# Import common utilities -from .generator import ( - COLDPRESS_LABELS, - generate_base_dir, - extract_container_config, -) - - -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 - - -# Use extract_container_config from common generator.py -_extract_container_config = extract_container_config - - -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/kubeflow_generator.py b/coldpress/kubeflow_generator.py index 448b7a1..db61e16 100644 --- a/coldpress/kubeflow_generator.py +++ b/coldpress/kubeflow_generator.py @@ -2,495 +2,13 @@ """Kubeflow manifest generation for Coldpress jobs.""" import yaml -from .generator import ( - COLDPRESS_LABELS, - generate_base_dir, - extract_container_config, - build_volumes_from_task, - build_volume_mounts_from_task, - detect_job_type, - detect_framework, +from .utils import ( + apply_arg_overrides, + apply_env_overrides, + build_discovery_init_container, ) -def generate_kubeflow(job_spec, node_assignments): - """Generate Kubeflow manifest (PyTorchJob or InferenceService). - - Auto-detects job type: - - Training job → PyTorchJob/TFJob/MPIJob - - Inference job → KServe InferenceService - - Args: - job_spec: Job specification dict - node_assignments: Node assignment dict - - Returns: - tuple: (manifest_dict, base_dir, manifest_type) - manifest_type is 'pytorchjob', 'tfjob', 'mpijob', or 'inferenceservice' - """ - tasks = job_spec["tasks"] - if len(tasks) == 0: - raise ValueError("No tasks specified in job spec") - - # For now, only handle single-task jobs - task = tasks[0] - job_type = detect_job_type(task) - - if job_type == "inference": - return generate_inferenceservice(job_spec, node_assignments) - else: - # Training job - framework = detect_framework(task) - if framework == "pytorch": - return generate_pytorchjob(job_spec, node_assignments) - elif framework == "tensorflow": - return generate_tfjob(job_spec, node_assignments) - elif framework == "mpi": - return generate_mpijob(job_spec, node_assignments) - else: - # Default to PyTorchJob if framework unclear - return generate_pytorchjob(job_spec, node_assignments) - - -def _build_mkdir_init_container(base_dir, storage_pvc_name): - """Build init container to create directory structure. - - Args: - base_dir: Base directory path to create - storage_pvc_name: Name of the PVC for storage - - Returns: - dict: Init container spec - """ - return { - "name": "mkdir", - "image": "registry.access.redhat.com/ubi9/ubi-minimal:latest", - "command": [ - "sh", - "-c", - f"mkdir -p /mnt/coldpress-data/{base_dir} && echo Created directory {base_dir}", - ], - "volumeMounts": [ - {"name": "coldpress-data", "mountPath": "/mnt/coldpress-data"} - ], - } - - -def generate_pytorchjob(job_spec, node_assignments): - """Generate PyTorchJob manifest from job spec. - - Args: - job_spec: Job specification dict - node_assignments: Node assignment dict - - Returns: - tuple: (pytorchjob_dict, base_dir, 'pytorchjob') - """ - namespace = job_spec["namespace"] - tasks = job_spec["tasks"] - task = tasks[0] # Single task for now - job_name = job_spec["name"] - - # Generate base directory for results - base_dir = generate_base_dir(namespace, job_name) - - # Get storage PVC name - storage = job_spec.get("storage", {}) - storage_pvc_name = storage.get("results", f"coldpress-{namespace}-storage") - - # Extract container configuration - container_config = extract_container_config(task) - - # Build container spec - container = { - "name": container_config.get("name", "pytorch"), - "image": container_config["image"], - } - - if container_config.get("command"): - container["command"] = container_config["command"] - if container_config.get("args"): - container["args"] = container_config["args"] - if container_config.get("working_dir"): - container["workingDir"] = container_config["working_dir"] - if container_config.get("env"): - container["env"] = container_config["env"] - - # Add resources - if container_config.get("resources"): - container["resources"] = container_config["resources"] - - # Build volume mounts - volume_mounts = build_volume_mounts_from_task(task) - - # Add ConfigMap volume mount if present - if "configmap" in job_spec: - volume_mounts.append({"name": "config-files", "mountPath": "/workspace"}) - - if volume_mounts: - container["volumeMounts"] = volume_mounts - - # Build volumes - volumes = build_volumes_from_task(task, job_spec) - - # Ensure coldpress-data PVC is in volumes (needed for init container) - if not any(v.get("name") == "coldpress-data" for v in volumes): - volumes.insert( - 0, - { - "name": "coldpress-data", - "persistentVolumeClaim": {"claimName": storage_pvc_name}, - }, - ) - - # Add ConfigMap volume if present - if "configmap" in job_spec: - cm_name = job_spec["configmap"]["name"] - # Add coldpress- prefix if not already present - if not cm_name.startswith("coldpress-"): - cm_name = f"coldpress-{cm_name}" - volumes.append({"name": "config-files", "configMap": {"name": cm_name}}) - - # Add init container to create directories - mkdir_init = _build_mkdir_init_container(base_dir, storage_pvc_name) - - # Build pod template - pod_template = { - "metadata": { - "annotations": { - "sidecar.istio.io/inject": "false" # Prescriptive: disable Istio - } - }, - "spec": { - "initContainers": [mkdir_init], - "containers": [container], - "restartPolicy": "OnFailure", - }, - } - - if volumes: - pod_template["spec"]["volumes"] = volumes - - # Apply node selector if specified - node_id = node_assignments.get(0, "any") - if node_id != "any": - pod_template["spec"]["nodeSelector"] = {"coldpress.node": str(node_id)} - - # Apply tolerations if task requires - if task.get("tolerate_all"): - pod_template["spec"]["tolerations"] = [{"operator": "Exists"}] - - # Build PyTorchJob manifest - pytorchjob_name = f"coldpress-{job_name}" - pytorchjob_labels = COLDPRESS_LABELS.copy() - pytorchjob_labels.update( - { - "coldpress.io/job-id": pytorchjob_name, - "kueue.x-k8s.io/queue-name": f"coldpress-local-queue-{namespace}", # Required for Kueue integration - } - ) - - # Add coldpress labels to pod template - if "labels" not in pod_template["metadata"]: - pod_template["metadata"]["labels"] = {} - pod_template["metadata"]["labels"].update(COLDPRESS_LABELS.copy()) - pod_template["metadata"]["labels"]["coldpress.io/job-id"] = pytorchjob_name - - pytorchjob = { - "apiVersion": "kubeflow.org/v1", - "kind": "PyTorchJob", - "metadata": { - "name": pytorchjob_name, - "namespace": namespace, - "labels": pytorchjob_labels, - "annotations": { - "coldpress.io/base-dir": base_dir, - "kueue.x-k8s.io/wait-for-pods-ready-timeout": "20m", # Allow time for image pulls - }, - }, - "spec": { - "pytorchReplicaSpecs": { - "Master": { - "replicas": 1, - "restartPolicy": "OnFailure", - "template": pod_template, - } - } - }, - } - - return pytorchjob, base_dir, "pytorchjob" - - -def generate_tfjob(job_spec, node_assignments): - """Generate TFJob manifest from job spec. - - Args: - job_spec: Job specification dict - node_assignments: Node assignment dict - - Returns: - tuple: (tfjob_dict, base_dir, 'tfjob') - """ - namespace = job_spec["namespace"] - tasks = job_spec["tasks"] - task = tasks[0] - job_name = job_spec["name"] - - base_dir = generate_base_dir(namespace, job_name) - - # Get storage PVC name - storage = job_spec.get("storage", {}) - storage_pvc_name = storage.get("results", f"coldpress-{namespace}-storage") - - container_config = extract_container_config(task) - - container = { - "name": "tensorflow", - "image": container_config["image"], - } - - if container_config.get("command"): - container["command"] = container_config["command"] - if container_config.get("args"): - container["args"] = container_config["args"] - if container_config.get("resources"): - container["resources"] = container_config["resources"] - - volume_mounts = build_volume_mounts_from_task(task) - if volume_mounts: - container["volumeMounts"] = volume_mounts - - volumes = build_volumes_from_task(task, job_spec) - - # Ensure coldpress-data PVC is in volumes (needed for init container) - if not any(v.get("name") == "coldpress-data" for v in volumes): - volumes.insert( - 0, - { - "name": "coldpress-data", - "persistentVolumeClaim": {"claimName": storage_pvc_name}, - }, - ) - - # Add init container to create directories - mkdir_init = _build_mkdir_init_container(base_dir, storage_pvc_name) - - pod_template = { - "metadata": {"annotations": {"sidecar.istio.io/inject": "false"}}, - "spec": { - "initContainers": [mkdir_init], - "containers": [container], - "restartPolicy": "OnFailure", - }, - } - - if volumes: - pod_template["spec"]["volumes"] = volumes - - # Add coldpress labels to pod template - if "labels" not in pod_template["metadata"]: - pod_template["metadata"]["labels"] = {} - tfjob_name = f"coldpress-{job_name}" - pod_template["metadata"]["labels"].update(COLDPRESS_LABELS.copy()) - pod_template["metadata"]["labels"]["coldpress.io/job-id"] = tfjob_name - - tfjob_labels = COLDPRESS_LABELS.copy() - tfjob_labels.update( - { - "coldpress.io/job-id": tfjob_name, - "kueue.x-k8s.io/queue-name": f"coldpress-local-queue-{namespace}", - } - ) - - tfjob = { - "apiVersion": "kubeflow.org/v1", - "kind": "TFJob", - "metadata": { - "name": tfjob_name, - "namespace": namespace, - "labels": tfjob_labels, - "annotations": { - "coldpress.io/base-dir": base_dir, - "kueue.x-k8s.io/wait-for-pods-ready-timeout": "20m", # Allow time for image pulls - }, - }, - "spec": { - "tfReplicaSpecs": { - "Worker": { - "replicas": 1, - "restartPolicy": "OnFailure", - "template": pod_template, - } - } - }, - } - - return tfjob, base_dir, "tfjob" - - -def generate_mpijob(job_spec, node_assignments): - """Generate MPIJob manifest from job spec. - - Args: - job_spec: Job specification dict - node_assignments: Node assignment dict - - Returns: - tuple: (mpijob_dict, base_dir, 'mpijob') - """ - namespace = job_spec["namespace"] - tasks = job_spec["tasks"] - task = tasks[0] - job_name = job_spec["name"] - - base_dir = generate_base_dir(namespace, job_name) - - # Get storage PVC name - storage = job_spec.get("storage", {}) - storage_pvc_name = storage.get("results", f"coldpress-{namespace}-storage") - - container_config = extract_container_config(task) - - launcher_container = { - "name": "mpi-launcher", - "image": container_config["image"], - } - - if container_config.get("command"): - launcher_container["command"] = container_config["command"] - if container_config.get("args"): - launcher_container["args"] = container_config["args"] - if container_config.get("resources"): - launcher_container["resources"] = container_config["resources"] - - worker_container = launcher_container.copy() - worker_container["name"] = "mpi-worker" - - mpijob_name = f"coldpress-{job_name}" - mpijob_labels = COLDPRESS_LABELS.copy() - mpijob_labels.update( - { - "coldpress.io/job-id": mpijob_name, - "kueue.x-k8s.io/queue-name": f"coldpress-local-queue-{namespace}", - } - ) - - # Add init container to create directories - mkdir_init = _build_mkdir_init_container(base_dir, storage_pvc_name) - - # Build pod templates with coldpress labels - launcher_pod_template = { - "metadata": { - "labels": COLDPRESS_LABELS.copy(), - "annotations": {"sidecar.istio.io/inject": "false"}, - }, - "spec": {"initContainers": [mkdir_init], "containers": [launcher_container]}, - } - launcher_pod_template["metadata"]["labels"]["coldpress.io/job-id"] = mpijob_name - - worker_pod_template = { - "metadata": { - "labels": COLDPRESS_LABELS.copy(), - "annotations": {"sidecar.istio.io/inject": "false"}, - }, - "spec": {"initContainers": [mkdir_init], "containers": [worker_container]}, - } - worker_pod_template["metadata"]["labels"]["coldpress.io/job-id"] = mpijob_name - - mpijob = { - "apiVersion": "kubeflow.org/v2beta1", - "kind": "MPIJob", - "metadata": { - "name": mpijob_name, - "namespace": namespace, - "labels": mpijob_labels, - "annotations": { - "coldpress.io/base-dir": base_dir, - "kueue.x-k8s.io/wait-for-pods-ready-timeout": "20m", # Allow time for image pulls - }, - }, - "spec": { - "slotsPerWorker": 1, - "runPolicy": {"cleanPodPolicy": "Running"}, - "mpiReplicaSpecs": { - "Launcher": {"replicas": 1, "template": launcher_pod_template}, - "Worker": {"replicas": 1, "template": worker_pod_template}, - }, - }, - } - - return mpijob, base_dir, "mpijob" - - -def generate_inferenceservice(job_spec, node_assignments): - """Generate KServe InferenceService manifest from job spec. - - Args: - job_spec: Job specification dict - node_assignments: Node assignment dict - - Returns: - tuple: (inferenceservice_dict, base_dir, 'inferenceservice') - """ - namespace = job_spec["namespace"] - tasks = job_spec["tasks"] - task = tasks[0] - job_name = job_spec["name"] - - base_dir = generate_base_dir(namespace, job_name) - - container_config = extract_container_config(task) - framework = detect_framework(task) - - # Build predictor based on framework - predictor = {} - - if framework == "pytorch": - predictor["pytorch"] = { - "storageUri": task.get("storageUri", "gs://bucket/model"), - "resources": container_config.get("resources", {}), - } - elif framework == "tensorflow": - predictor["tensorflow"] = { - "storageUri": task.get("storageUri", "gs://bucket/model"), - "resources": container_config.get("resources", {}), - } - else: - # Use custom container - predictor["containers"] = [ - { - "name": "kserve-container", - "image": container_config["image"], - "resources": container_config.get("resources", {}), - } - ] - - inferenceservice_name = f"coldpress-{job_name}" - inferenceservice_labels = COLDPRESS_LABELS.copy() - inferenceservice_labels.update( - { - "coldpress.io/job-id": inferenceservice_name, - } - ) - - inferenceservice = { - "apiVersion": "serving.kserve.io/v1beta1", - "kind": "InferenceService", - "metadata": { - "name": inferenceservice_name, - "namespace": namespace, - "labels": inferenceservice_labels, - "annotations": { - "coldpress.io/base-dir": base_dir, - }, - }, - "spec": {"predictor": predictor}, - } - - return inferenceservice, base_dir, "inferenceservice" - - def generate_inferenceservice_from_intent( vk8s_jobs: dict, intent_config, project_config, namespace: str ): @@ -507,7 +25,6 @@ def generate_inferenceservice_from_intent( tuple: (InferenceService manifest, base_dir) """ import copy - import re from datetime import datetime, timezone from .constants import ( COLDPRESS_LABELS, @@ -530,8 +47,12 @@ def generate_inferenceservice_from_intent( job_id = task_name inferenceservice_name = f"coldpress-{job_id}" - storage = project_config.get("storage", {}) - data_pvc_name = storage.get("results", f"coldpress-{namespace}-storage") + # Storage is validated by Pydantic - if present, 'results' is required + storage = project_config.get("storage") + if storage: + data_pvc_name = storage["results"] + else: + data_pvc_name = f"coldpress-{namespace}-storage" # Generate base directory for results/logs timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") @@ -555,88 +76,15 @@ def generate_inferenceservice_from_intent( # Get replicas (for autoscaling configuration) replicas = task_intent.replicas or 1 - # Build macros for arg/env substitution - def substitute_macros(value: str, macros: dict) -> str: - """Substitute ${MACRO} patterns in a string.""" - if not isinstance(value, str): - return value - - pattern = r"\$\{([^}]+)\}" - - def replacer(match): - macro_name = match.group(1) - if macro_name not in macros: - return match.group(0) - return str(macros[macro_name]) - - return re.sub(pattern, replacer, value) - # KServe macros macros = { "REPLICAS": str(replicas), "TASK_NAME": task_name, } - # Apply arg replacements - if task_intent.args and container.get("args"): - args = ( - container["args"] - if isinstance(container["args"], list) - else [container["args"]] - ) - - for arg_key, arg_value in task_intent.args.items(): - substituted_value = substitute_macros(arg_value, macros) - - # Replace or add argument - new_args = [] - i = 0 - found = False - - while i < len(args): - arg = str(args[i]) - - if arg.startswith(f"--{arg_key}="): - new_args.append(f"--{arg_key}={substituted_value}") - found = True - i += 1 - elif arg == f"--{arg_key}": - new_args.append(arg) - new_args.append(substituted_value) - found = True - i += 2 - else: - new_args.append(arg) - i += 1 - - if not found: - new_args.append(f"--{arg_key}={substituted_value}") - - args = new_args - - container["args"] = args - - # Apply env var additions/replacements - if task_intent.env: - env_list = container.get("env", []) - if not isinstance(env_list, list): - env_list = [] - - for env_key, env_value in task_intent.env.items(): - substituted_value = substitute_macros(str(env_value), macros) - - # 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 + # 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", []) @@ -710,73 +158,7 @@ def kubeflow_to_yaml(manifest): Returns: str: YAML string """ - return yaml.dump(manifest, default_flow_style=False, sort_keys=False) - - -def _build_discovery_init_container( - template_path: str, base_dir: str, storage_volume_name: str -): - """ - Build discovery init container from template for Kubeflow. - - Args: - template_path: Path to discovery template YAML - base_dir: Base directory path - storage_volume_name: Name of the storage volume to mount - - Returns: - dict: Init container spec, or None if template not found - """ - import os - import sys - import yaml - import copy - - 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 = copy.deepcopy(containers[0]) - container["name"] = "discovery" - - # Update volume mounts to use PVC with base directory - container["volumeMounts"] = [ - { - "name": storage_volume_name, - "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] - - 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 + return yaml.safe_dump(manifest, default_flow_style=False, sort_keys=False) def generate_pytorchjob_from_intent( @@ -796,11 +178,11 @@ def generate_pytorchjob_from_intent( """ import copy import os - import re from datetime import datetime, timezone from .constants import ( COLDPRESS_LABELS, MKDIR_IMAGE, + DEFAULT_DISCOVERY_DIR, get_kueue_queue_label, ) @@ -820,8 +202,12 @@ def generate_pytorchjob_from_intent( job_id = task_name pytorchjob_name = f"coldpress-{job_id}" - storage = project_config.get("storage", {}) - data_pvc_name = storage.get("results", f"coldpress-{namespace}-storage") + # Storage is validated by Pydantic - if present, 'results' is required + storage = project_config.get("storage") + if storage: + data_pvc_name = storage["results"] + else: + data_pvc_name = f"coldpress-{namespace}-storage" # Generate base directory for results timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") @@ -841,7 +227,7 @@ def generate_pytorchjob_from_intent( ) # Resolve template path - discovery_dir = os.getenv("COLDPRESS_DISCOVERY_DIR", "discovery") + 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 @@ -871,88 +257,15 @@ def generate_pytorchjob_from_intent( # Get replicas (PyTorchJob workers) replicas = task_intent.replicas or 1 - # Build macros for arg substitution (using Kubeflow auto-injected vars) - def substitute_macros(value: str, macros: dict) -> str: - """Substitute ${MACRO} patterns in a string.""" - if not isinstance(value, str): - return value - - pattern = r"\$\{([^}]+)\}" - - def replacer(match): - macro_name = match.group(1) - if macro_name not in macros: - return match.group(0) - return str(macros[macro_name]) - - return re.sub(pattern, replacer, value) - # Kubeflow macros - simpler because PyTorchJob auto-injects env vars macros = { "REPLICAS": str(replicas), "TASK_NAME": task_name, } - # Apply arg replacements - if task_intent.args and container.get("args"): - args = ( - container["args"] - if isinstance(container["args"], list) - else [container["args"]] - ) - - for arg_key, arg_value in task_intent.args.items(): - substituted_value = substitute_macros(arg_value, macros) - - # Replace or add argument - new_args = [] - i = 0 - found = False - - while i < len(args): - arg = str(args[i]) - - if arg.startswith(f"--{arg_key}="): - new_args.append(f"--{arg_key}={substituted_value}") - found = True - i += 1 - elif arg == f"--{arg_key}": - new_args.append(arg) - new_args.append(substituted_value) - found = True - i += 2 - else: - new_args.append(arg) - i += 1 - - if not found: - new_args.append(f"--{arg_key}={substituted_value}") - - args = new_args - - container["args"] = args - - # Apply env var additions/replacements - if task_intent.env: - env_list = container.get("env", []) - if not isinstance(env_list, list): - env_list = [] - - for env_key, env_value in task_intent.env.items(): - substituted_value = substitute_macros(str(env_value), macros) - - # 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 + # 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 = { @@ -1012,7 +325,7 @@ def replacer(match): break if storage_volume_name: - discovery_init = _build_discovery_init_container( + discovery_init = build_discovery_init_container( discovery_template_path, base_dir, storage_volume_name ) if discovery_init: diff --git a/coldpress/kuberay_generator.py b/coldpress/kuberay_generator.py new file mode 100644 index 0000000..6fb59c0 --- /dev/null +++ b/coldpress/kuberay_generator.py @@ -0,0 +1,511 @@ +# Generated by: Claude Sonnet 4.5 +"""KubeRay manifest generation for Coldpress jobs. + +Transforms single-node Ray Train jobs into distributed RayJob manifests. +""" + +import os +import copy +import yaml +from datetime import datetime, timezone +from .constants import ( + COLDPRESS_LABELS, + MKDIR_IMAGE, + DEFAULT_DISCOVERY_DIR, + get_kueue_queue_label, +) +from .utils import ( + apply_arg_overrides, + apply_env_overrides, + build_discovery_init_container, +) + + +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}" + + # Storage is validated by Pydantic - if present, 'results' is required + storage = project_config.get("storage") + if storage: + data_pvc_name = storage["results"] + else: + data_pvc_name = f"coldpress-{namespace}-storage" + + # 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 = 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_task_names = set(discovery_tasks) + + # 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 = [] + + # 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) + + # 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 9d87b62..5b65b83 100644 --- a/coldpress/script_gen.py +++ b/coldpress/script_gen.py @@ -1,13 +1,111 @@ # Generated 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", + timeout=None, configmap_name=None, configmap_files=None, manifest_type="jobset", @@ -18,7 +116,7 @@ def generate_run_script( Args: 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) @@ -26,9 +124,25 @@ def generate_run_script( 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}" @@ -40,50 +154,32 @@ def generate_run_script( oc create configmap {actual_configmap_name} -n {namespace} {from_files} --dry-run=client -o yaml | oc apply -f - """ - # Generate manifest-specific commands - if manifest_type == "jobset": - manifest_file = "jobset.yaml" - resource_type = "jobset" - resource_name = f"coldpress-{job_name}" - apply_cmd = 'echo "Applying JobSet..."\noc apply -f jobset.yaml' - services_apply = """ + # 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""" - elif manifest_type == "pytorchjob": - manifest_file = "pytorchjob.yaml" - resource_type = "pytorchjob" - resource_name = f"coldpress-{job_name}" - apply_cmd = 'echo "Applying PyTorchJob..."\noc apply -f pytorchjob.yaml' - services_apply = "" - elif manifest_type == "tfjob": - manifest_file = "tfjob.yaml" - resource_type = "tfjob" - resource_name = f"coldpress-{job_name}" - apply_cmd = 'echo "Applying TFJob..."\noc apply -f tfjob.yaml' - services_apply = "" - elif manifest_type == "mpijob": - manifest_file = "mpijob.yaml" - resource_type = "mpijob" - resource_name = f"coldpress-{job_name}" - apply_cmd = 'echo "Applying MPIJob..."\noc apply -f mpijob.yaml' - services_apply = "" - elif manifest_type == "inferenceservice": - manifest_file = "kservejob.yaml" - resource_type = "inferenceservice" - resource_name = f"coldpress-{job_name}" - apply_cmd = ( - 'echo "Applying KServe InferenceService..."\noc apply -f kservejob.yaml' - ) - services_apply = "" - else: - manifest_file = "manifest.yaml" - resource_type = manifest_type - resource_name = f"coldpress-{job_name}" - apply_cmd = f'echo "Applying {manifest_type}..."\noc apply -f {manifest_file}' - services_apply = "" + if config.get("has_services", False) + else "" + ) # Determine wait condition based on resource type if manifest_type == "inferenceservice": @@ -146,6 +242,10 @@ def generate_cleanup_script( 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 @@ -229,6 +329,10 @@ def generate_monitor_script(job_name, namespace, manifest_type="jobset"): 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" @@ -284,6 +388,11 @@ 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 @@ -335,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 @@ -395,6 +504,11 @@ 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 @@ -427,8 +541,8 @@ def generate_explore_script(job_name, namespace, storage_pvc, base_dir): 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 @@ -491,6 +605,11 @@ 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 @@ -528,8 +647,8 @@ def generate_copy_script(job_name, namespace, storage_pvc, base_dir): 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 diff --git a/coldpress/utils.py b/coldpress/utils.py new file mode 100644 index 0000000..d712c2c --- /dev/null +++ b/coldpress/utils.py @@ -0,0 +1,367 @@ +# Generated 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 + + +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 + dangerous_chars = ["`", "$", "|", ";", "&", ">", "<", "\n", "\r"] + for char in dangerous_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 build_discovery_init_container( + template_path: str, + base_dir: str, + storage_volume_name: str, + task_id: int | None = None, +): + """ + 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) + + 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 = copy.deepcopy(containers[0]) + container["name"] = "discovery" + + # Construct result path + if task_id is not None: + # JobSet mode - task-specific path + result_path = f"{base_dir}/task-{task_id}" + else: + # KubeRay/Kubeflow mode - base directory only + result_path = base_dir + + # Update volume mounts to use PVC + container["volumeMounts"] = [ + { + "name": storage_volume_name, + "mountPath": "/tmp/result", + "subPath": 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 diff --git a/coldpress_common/__init__.py b/coldpress_common/__init__.py index 491deec..56678e2 100644 --- a/coldpress_common/__init__.py +++ b/coldpress_common/__init__.py @@ -2,11 +2,9 @@ """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, @@ -22,11 +20,9 @@ ) __all__ = [ - "validate_config", "validate_intent", "validate_project_config", "validate_task_specs", - "validate_job_spec", "validate_user_config", "validate_cluster_config", "validate_kubernetes_name", diff --git a/coldpress_common/model.py b/coldpress_common/model.py index d4d71cc..807c4ed 100644 --- a/coldpress_common/model.py +++ b/coldpress_common/model.py @@ -40,7 +40,14 @@ """ 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 === @@ -249,13 +256,32 @@ class DependsOnConfig(BaseModel): 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, str]] = None # Key-value pairs for arg replacement + 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 @@ -281,7 +307,9 @@ class IntentConfig(BaseModel): project: str output: str - target: Literal["jobset", "kubeflow", "kserve"] = "jobset" # Backend to generate + target: Literal["jobset", "kubeflow", "kserve", "kuberay"] = ( + "jobset" # Backend to generate + ) files: Optional[list[str]] = None discovery: Optional[Union[str, DiscoveryConfig]] = None tasks: list[TaskIntent] @@ -476,22 +504,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. @@ -525,28 +537,12 @@ 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: - """ - Validate complete job specification. - - Args: - job_spec_data: Complete job spec dict - - Returns: - Validated JobSpec - - Raises: - ValidationError: If validation fails - """ - return JobSpec.model_validate(job_spec_data) - - def validate_user_config(user_data: dict) -> UserConfig: """ Validate user configuration. diff --git a/coldpress_setup/cli.py b/coldpress_setup/cli.py index 928f5f4..86708f4 100644 --- a/coldpress_setup/cli.py +++ b/coldpress_setup/cli.py @@ -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 @@ -225,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) @@ -241,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 diff --git a/coldpress_setup/generator.py b/coldpress_setup/generator.py index 3d718b0..85e454e 100644 --- a/coldpress_setup/generator.py +++ b/coldpress_setup/generator.py @@ -1,14 +1,8 @@ # Generated 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}", @@ -168,6 +153,7 @@ def generate_rbac(namespace, targets="jobset"): 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 = [] @@ -200,6 +186,16 @@ def generate_rbac(namespace, targets="jobset"): } ) + # 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( [ @@ -331,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). @@ -495,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 @@ -510,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/examples/pytorch_ddp_training/ARG_OVERRIDE_README.md b/examples/pytorch_ddp_training/ARG_OVERRIDE_README.md new file mode 100644 index 0000000..1a829c0 --- /dev/null +++ b/examples/pytorch_ddp_training/ARG_OVERRIDE_README.md @@ -0,0 +1,116 @@ +# Argument Override Pattern for Coldpress + +This example demonstrates the new argument override pattern for Coldpress intent files. + +## Problem + +The previous approach used environment variables with complex bash conditionals in the job-spec, making YAMLs unreadable and difficult to maintain. Single-node job-specs couldn't be directly run with `kubectl apply`. + +## Solution + +Use clean, directly-runnable job-specs with argument substitution from intent.yaml. + +## Key Features + +### 1. Clean Single-Node Job-Spec + +The `job-spec.yaml` is now a standard Kubernetes Job that can be directly applied: + +```yaml +command: + - python + - -m + - torch.distributed.run + - --nproc_per_node=2 + - train.py + - --dataset=mnist + - --epochs=50 + - --batch-size=128 +``` + +### 2. Intent-Based Argument Overrides + +Two patterns are supported in `intent.yaml`: + +#### Simple Replacement (existing args) +```yaml +args: + epochs: "100" # Replace existing --epochs value + batch-size: "256" # Replace existing --batch-size value +``` + +#### Positioned Insertion (new args) +```yaml +args: + nnodes: + value: "${REPLICAS}" + insert_after: nproc_per_node + + master_addr: + value: "${REPLICA_ddp-training_0}" + insert_after: node_rank +``` + +## Generated Output + +The JobSet generator produces: + +```yaml +command: + - python + - -m + - torch.distributed.run + - --nproc_per_node=2 + - --nnodes=2 # Inserted after nproc_per_node + - --node_rank=0 # Inserted after nnodes + - --master_addr=coldpress-... # Inserted after node_rank + - --master_port=29500 # Inserted after master_addr + - train.py + - --dataset=mnist + - --epochs=100 # Replaced from 50 + - --batch-size=256 # Replaced from 128 +``` + +## Macro Substitution + +The following macros are available: +- `${REPLICAS}` - Total replica count +- `${INDEX}` - Current replica index (0-based) +- `${REPLICA__N}` - DNS name of replica N +- `${NODE_ID}` - Assigned node ID (if using node pinning) + +## Implementation + +### Model Changes + +`coldpress_common/model.py`: +- Added `ArgOverride` Pydantic model +- Updated `TaskIntent.args` to support `Union[str, ArgOverride]` + +### Generator Changes + +Both `jobset_generator.py` and `kubeflow_generator.py`: +- Added `parse_arg()` function to parse CLI flags +- Enhanced `replace_or_add_arg()` with `insert_after`/`insert_before` support +- Updated arg processing to handle both `command` and `args` fields +- Handle both string values and ArgOverride objects + +## Example Files + +- `job-spec.yaml` - Clean single-node Job specification +- `intent_jobset.yaml` - JobSet intent with arg overrides +- `intent_kubeflow.yaml` - PyTorchJob intent (simpler, uses auto-injection) + +## Testing + +Generate manifests: +```bash +coldpress generate --intent intent_jobset.yaml +coldpress generate --intent intent_kubeflow.yaml +``` + +Verify args in output: +```bash +grep -A 20 "command:" output/ddp-training-job/jobset.yaml +grep -A 20 "command:" output/ddp-training-job/pytorchjob.yaml +``` diff --git a/examples/pytorch_ddp_training/intent_jobset.yaml b/examples/pytorch_ddp_training/intent_jobset.yaml index ddfbd6c..92fcaec 100644 --- a/examples/pytorch_ddp_training/intent_jobset.yaml +++ b/examples/pytorch_ddp_training/intent_jobset.yaml @@ -12,10 +12,22 @@ discovery: tasks: - name: ddp-training replicas: 2 - env: - # Don't override NPROC_PER_NODE - use value from job-spec (2 GPUs per node) - NNODES: "${REPLICAS}" - MASTER_ADDR: "${REPLICA_ddp-training_0}" - MASTER_PORT: "29500" - RANK: "${INDEX}" - WORLD_SIZE: "${REPLICAS}" + + 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 index a4341ab..a574856 100644 --- a/examples/pytorch_ddp_training/intent_kubeflow.yaml +++ b/examples/pytorch_ddp_training/intent_kubeflow.yaml @@ -12,10 +12,10 @@ discovery: tasks: - name: ddp-training replicas: 2 + args: - nnodes: "${REPLICAS}" - # Kubeflow PyTorchJob auto-injects: - # - MASTER_ADDR - # - MASTER_PORT - # - RANK (via RANK env var) - # So we don't need to specify them here + # 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 7d3229c..9864052 100644 --- a/examples/pytorch_ddp_training/job-spec.yaml +++ b/examples/pytorch_ddp_training/job-spec.yaml @@ -16,31 +16,19 @@ spec: - name: training image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime workingDir: /workspace - command: ["bash", "-c"] - args: - - | - # Build torch.distributed.run command from env vars - CMD="python -m torch.distributed.run" - - # Add distributed params if env vars are set - [ -n "$NPROC_PER_NODE" ] && CMD="$CMD --nproc_per_node=$NPROC_PER_NODE" - [ -n "$NNODES" ] && CMD="$CMD --nnodes=$NNODES" - [ -n "$RANK" ] && CMD="$CMD --node_rank=$RANK" - [ -n "$MASTER_ADDR" ] && CMD="$CMD --master_addr=$MASTER_ADDR" - [ -n "$MASTER_PORT" ] && CMD="$CMD --master_port=$MASTER_PORT" - - # Training script and args - CMD="$CMD train.py" - CMD="$CMD --dataset=mnist" - CMD="$CMD --train-test-split=0.8" - CMD="$CMD --epochs=50" - CMD="$CMD --batch-size=128" - CMD="$CMD --hidden-size=4096" - CMD="$CMD --lr=0.01" - CMD="$CMD --output-dir=/results/checkpoints" - - echo "Running: $CMD" - exec $CMD + 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" @@ -52,8 +40,6 @@ spec: env: - name: NCCL_DEBUG value: "INFO" - - name: NPROC_PER_NODE - value: "2" volumeMounts: - name: results mountPath: /results diff --git a/examples/pytorch_ray_training/ARG_OVERRIDE_README.md b/examples/pytorch_ray_training/ARG_OVERRIDE_README.md new file mode 100644 index 0000000..884728b --- /dev/null +++ b/examples/pytorch_ray_training/ARG_OVERRIDE_README.md @@ -0,0 +1,105 @@ +# Argument Override Pattern for Ray Training Example + +This example demonstrates the argument override pattern for Ray Train jobs using KubeRay, matching the approach used in the PyTorch DDP example. + +## Updates Made + +### 1. Clean Single-Node Job-Spec + +The `job-spec.yaml` now uses clean command/args (no bash conditionals): + +```yaml +command: + - python + - train.py + - --dataset=mnist + - --epochs=50 + - --batch-size=128 + - --hidden-size=4096 + - --lr=0.01 + - --num-workers=1 # Single-node default + - --gpus-per-worker=2 + - --cpus-per-worker=4 +``` + +**Before:** Used bash script with env var conditionals (complex, hard to read) +**After:** Direct command with arguments (clean, directly runnable) + +### 2. Intent-Based Argument Overrides + +The `intent_kuberay.yaml` now uses `args` instead of `env`: + +```yaml +tasks: + - name: ray-training + replicas: 2 # 1 head + 1 worker = 4 GPUs total + + args: + # Scale Ray Train workers to match pod count + num-workers: "${REPLICAS}" # Overrides from 1 to 2 + # gpus-per-worker and cpus-per-worker use defaults from job-spec +``` + +**Before:** Used `env` with environment variable overrides +**After:** Uses `args` with macro substitution (consistent with DDP/Kubeflow patterns) + +### 3. Generator Updates + +Updated `kuberay_generator.py` to: +- Import `parse_arg` and `replace_or_add_arg` from `jobset_generator` +- Apply arg overrides before building the RayJob entrypoint +- Handle both simple string values and `ArgOverride` objects with positioning +- Process args for both head and worker containers + +## Generated Output + +The RayJob manifest now has a clean entrypoint: + +```yaml +entrypoint: cd /results/.../workspace && 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=2 # Overridden from 1 + --gpus-per-worker=2 # Default preserved + --cpus-per-worker=4 # Default preserved +``` + +## Macro Substitution + +Available macros for RayJob: +- `${REPLICAS}` - Total replica count (head + workers) +- `${INDEX}` - Container index (0 for head, 1+ for workers) +- `${TASK_NAME}` - Task name from intent + +## Consistency with DDP Example + +Both examples now follow the same pattern: + +| Aspect | DDP Example | Ray Example | +|--------|-------------|-------------| +| job-spec | Clean command/args | Clean command/args | +| intent | Args with positioning | Args with macros | +| Generator | `jobset_generator.py` | `kuberay_generator.py` | +| Override function | `replace_or_add_arg()` | `replace_or_add_arg()` | + +## Testing + +Generate and verify: +```bash +coldpress generate --intent examples/pytorch_ray_training/intent_kuberay.yaml +grep "num-workers" output/ray-training-job/rayjob.yaml +# Should show: --num-workers=2 (overridden from 1) +``` + +## Benefits + +1. **Readable job-specs**: No bash scripts, directly runnable +2. **Consistent pattern**: Same approach across JobSet, Kubeflow, KubeRay +3. **Type-safe**: Pydantic validation of arg overrides +4. **Flexible**: Supports both simple replacement and positioned insertion +5. **Maintainable**: Clear separation between single-node defaults and multi-node scaling 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/intent_jobset.yaml b/examples/vllm_guidellm_benchmark/intent_jobset.yaml index e59aa85..b8c6da2 100644 --- a/examples/vllm_guidellm_benchmark/intent_jobset.yaml +++ b/examples/vllm_guidellm_benchmark/intent_jobset.yaml @@ -15,5 +15,5 @@ tasks: depends_on: task: inference-server wait_for: ready - env: - GUIDELLM_TARGET: "http://${REPLICA_inference-server_0}:8000" + args: + target: "http://${REPLICA_inference-server_0}:8000" diff --git a/examples/vllm_guidellm_benchmark/job-spec.yaml b/examples/vllm_guidellm_benchmark/job-spec.yaml index 74a687e..dbca378 100644 --- a/examples/vllm_guidellm_benchmark/job-spec.yaml +++ b/examples/vllm_guidellm_benchmark/job-spec.yaml @@ -21,25 +21,17 @@ spec: 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}" + 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 - - 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 @@ -80,34 +72,20 @@ spec: 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}" + 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 - - name: GUIDELLM_TARGET - value: "http://localhost: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 diff --git a/projects/coldpress-project.yaml b/projects/coldpress-project.yaml index ba6da1e..76fede9 100644 --- a/projects/coldpress-project.yaml +++ b/projects/coldpress-project.yaml @@ -2,8 +2,8 @@ namespace: coldpress-project # Target platforms - determines which RBAC permissions to grant -# Options: jobset, kubeflow, or both (comma-separated) -targets: jobset,kubeflow +# Options: jobset, kubeflow, kuberay, or combination (comma-separated) +targets: jobset,kubeflow,kuberay # References to cluster-wide resources cluster_queue: coldpress-cluster-queue diff --git a/pyproject.toml b/pyproject.toml index 82e0507..86a4de3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ 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" diff --git a/tests/test_labels.py b/tests/test_labels.py index 91f3983..a897f09 100644 --- a/tests/test_labels.py +++ b/tests/test_labels.py @@ -25,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, ( @@ -78,7 +78,7 @@ def test_jobset_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 labels["app.kubernetes.io/version"] == "0.2.1" print(f"✅ JobSet labels: {labels}") @@ -233,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: From 1208c59652b8748555b8867c778f3fdc246526fe Mon Sep 17 00:00:00 2001 From: asanaullah Date: Wed, 29 Apr 2026 22:43:56 -0700 Subject: [PATCH 03/14] applied vibe coding guidelines --- coldpress/cli.py | 92 +++++------------------ coldpress/constants.py | 12 +++ coldpress/jobset_generator.py | 50 ++++--------- coldpress/kubeflow_generator.py | 77 ++++++------------- coldpress/kuberay_generator.py | 50 +++++-------- coldpress/utils.py | 127 +++++++++++++++++++++++++++++++- coldpress_common/model.py | 3 +- 7 files changed, 213 insertions(+), 198 deletions(-) diff --git a/coldpress/cli.py b/coldpress/cli.py index 167b756..f8cbe20 100644 --- a/coldpress/cli.py +++ b/coldpress/cli.py @@ -16,6 +16,8 @@ import yaml import os import json +import shutil +import traceback from datetime import datetime, timezone from .jobset_generator import ( @@ -33,18 +35,19 @@ kuberay_to_yaml, ) from .script_gen import write_scripts -from .constants import COLDPRESS_VERSION +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_intent, validate_project_config, ) 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=COLDPRESS_VERSION) @@ -103,7 +106,9 @@ def _load_vk8s_jobs(job_spec_path): 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}") @@ -152,6 +157,9 @@ def generate(intent): coldpress generate -i examples/vllm_guidellm_benchmark/intent.yaml """ try: + # Read config at call time to enable testing and runtime reconfiguration + output_dir = os.getenv("COLDPRESS_OUTPUT_DIR", DEFAULT_OUTPUT_DIR) + # Load and validate intent.yaml intent_config = _load_intent(intent) intent_dir = os.path.dirname(os.path.abspath(intent)) @@ -164,7 +172,7 @@ def generate(intent): project_config, namespace = _load_project_config(intent_config.project) # Output path - output_path = os.path.join(OUTPUT_DIR, intent_config.output) + output_path = os.path.join(output_dir, intent_config.output) # Display generation info click.echo("Generating manifests") @@ -178,13 +186,7 @@ def generate(intent): # Use backend from intent.yaml backend = intent_config.target - backend_name_map = { - "kubeflow": "Kubeflow", - "kserve": "KServe", - "kuberay": "KubeRay", - "jobset": "JobSet", - } - backend_name = backend_name_map.get(backend, "JobSet") + backend_name = BACKEND_NAME_MAP.get(backend, "JobSet") click.echo(f"Target: {backend_name}") # Generate manifests based on target @@ -247,8 +249,6 @@ def generate(intent): # Copy ConfigMap files if intent_config.files: - import shutil - 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) @@ -275,65 +275,15 @@ def generate(intent): configmap_name = None configmap_files = [] if intent_config.files: - # Extract ConfigMap name from the manifest - 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: - configmap_name = volume["configMap"]["name"] - break - if configmap_name: - break - elif manifest_type == "pytorchjob": - # Search for configMap volumes in PyTorchJob manifest - replica_specs = manifest.get("spec", {}).get("pytorchReplicaSpecs", {}) - for replica_type, replica_spec in replica_specs.items(): - volumes = ( - replica_spec.get("template", {}) - .get("spec", {}) - .get("volumes", []) - ) - for volume in volumes: - if "configMap" in volume: - configmap_name = volume["configMap"]["name"] - break - if configmap_name: - break - 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: - configmap_name = volume["configMap"]["name"] - break + # 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 is validated by Pydantic - if present, 'results' is required - storage = project_config.get("storage") - if storage: - storage_pvc = storage["results"] - else: - storage_pvc = f"coldpress-{namespace}-storage" + 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-"): @@ -383,8 +333,6 @@ def generate(intent): raise SystemExit(1) 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 index eaf003e..68533a2 100644 --- a/coldpress/constants.py +++ b/coldpress/constants.py @@ -84,3 +84,15 @@ def get_pvc_name(namespace: str) -> str: "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/jobset_generator.py b/coldpress/jobset_generator.py index 9312dcf..5ef9c2a 100644 --- a/coldpress/jobset_generator.py +++ b/coldpress/jobset_generator.py @@ -9,22 +9,22 @@ """ import copy -import os import yaml from datetime import datetime, timezone from .constants import ( COLDPRESS_LABELS, MKDIR_IMAGE, - DEFAULT_DISCOVERY_DIR, get_kueue_queue_label, get_jobset_name, - get_pvc_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, ) @@ -83,12 +83,7 @@ def generate_jobset_from_intent( job_id = f"{intent_config.tasks[0].name}-workflow" jobset_name = get_jobset_name(job_id) - # Storage is validated by Pydantic - if present, 'results' is required - storage = project_config.get("storage") - if storage: - data_pvc_name = storage["results"] - else: - data_pvc_name = get_pvc_name(namespace) + 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") @@ -96,31 +91,9 @@ def generate_jobset_from_intent( base_dir = f"{namespace}/coldpress_results/{job_id}-{short_hash:08x}-{timestamp}" # Extract discovery configuration - 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) + discovery_template_path, discovery_task_names = parse_discovery_config( + intent_config + ) # Build replicated jobs replicated_jobs = [] @@ -286,6 +259,15 @@ 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)) diff --git a/coldpress/kubeflow_generator.py b/coldpress/kubeflow_generator.py index db61e16..69235c5 100644 --- a/coldpress/kubeflow_generator.py +++ b/coldpress/kubeflow_generator.py @@ -1,11 +1,22 @@ # Generated 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, ) @@ -24,13 +35,6 @@ def generate_inferenceservice_from_intent( Returns: tuple: (InferenceService manifest, base_dir) """ - import copy - from datetime import datetime, timezone - from .constants import ( - COLDPRESS_LABELS, - get_kueue_queue_label, - ) - # Only support single task for KServe if len(intent_config.tasks) != 1: raise ValueError("KServe generation currently only supports single task") @@ -47,12 +51,7 @@ def generate_inferenceservice_from_intent( job_id = task_name inferenceservice_name = f"coldpress-{job_id}" - # Storage is validated by Pydantic - if present, 'results' is required - storage = project_config.get("storage") - if storage: - data_pvc_name = storage["results"] - else: - data_pvc_name = f"coldpress-{namespace}-storage" + 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") @@ -176,16 +175,6 @@ def generate_pytorchjob_from_intent( Returns: tuple: (PyTorchJob manifest, base_dir) """ - import copy - import os - from datetime import datetime, timezone - from .constants import ( - COLDPRESS_LABELS, - MKDIR_IMAGE, - DEFAULT_DISCOVERY_DIR, - get_kueue_queue_label, - ) - # Only support single task for Kubeflow if len(intent_config.tasks) != 1: raise ValueError("Kubeflow generation currently only supports single task") @@ -202,12 +191,7 @@ def generate_pytorchjob_from_intent( job_id = task_name pytorchjob_name = f"coldpress-{job_id}" - # Storage is validated by Pydantic - if present, 'results' is required - storage = project_config.get("storage") - if storage: - data_pvc_name = storage["results"] - else: - data_pvc_name = f"coldpress-{namespace}-storage" + 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") @@ -215,30 +199,9 @@ def generate_pytorchjob_from_intent( base_dir = f"{namespace}/coldpress_results/{job_id}-{short_hash:08x}-{timestamp}" # Extract discovery configuration - 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_task_names = set(discovery_tasks) + 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", {}) @@ -299,6 +262,14 @@ def generate_pytorchjob_from_intent( # 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", diff --git a/coldpress/kuberay_generator.py b/coldpress/kuberay_generator.py index 6fb59c0..e699a3e 100644 --- a/coldpress/kuberay_generator.py +++ b/coldpress/kuberay_generator.py @@ -4,20 +4,21 @@ Transforms single-node Ray Train jobs into distributed RayJob manifests. """ -import os import copy import yaml from datetime import datetime, timezone from .constants import ( COLDPRESS_LABELS, MKDIR_IMAGE, - DEFAULT_DISCOVERY_DIR, 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, ) @@ -57,13 +58,7 @@ def generate_rayjob_from_intent( # Determine job name job_id = task_name rayjob_name = f"coldpress-{job_id}" - - # Storage is validated by Pydantic - if present, 'results' is required - storage = project_config.get("storage") - if storage: - data_pvc_name = storage["results"] - else: - data_pvc_name = f"coldpress-{namespace}-storage" + 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") @@ -71,30 +66,9 @@ def generate_rayjob_from_intent( base_dir = f"{namespace}/coldpress_results/{job_id}-{short_hash:08x}-{timestamp}" # Extract discovery configuration - 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_task_names = set(discovery_tasks) + 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", {}) @@ -197,6 +171,15 @@ def substitute_output_paths(args_list): # 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 = ( @@ -216,6 +199,7 @@ def substitute_output_paths(args_list): } 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) diff --git a/coldpress/utils.py b/coldpress/utils.py index d712c2c..82e9f27 100644 --- a/coldpress/utils.py +++ b/coldpress/utils.py @@ -11,6 +11,9 @@ import sys import yaml +# Security: Dangerous shell metacharacters that could enable injection +DANGEROUS_SHELL_CHARS = ["`", "$", "|", ";", "&", ">", "<", "\n", "\r"] + def sanitize_value(value: str, context: str = "") -> str: """ @@ -30,8 +33,7 @@ def sanitize_value(value: str, context: str = "") -> str: return str(value) # Check for shell metacharacters that could be dangerous in args - dangerous_chars = ["`", "$", "|", ";", "&", ">", "<", "\n", "\r"] - for char in dangerous_chars: + for char in DANGEROUS_SHELL_CHARS: if char in value: raise ValueError( f"Invalid character '{char}' in value '{value}'{context}. " @@ -294,6 +296,118 @@ def apply_env_overrides(container: dict, task_intent, macros: dict[str, str]) -> 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, @@ -324,10 +438,14 @@ def build_discovery_init_container( if not containers: return None - # Get template name from filename + # 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]) @@ -351,10 +469,11 @@ def build_discovery_init_container( ] # 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 /tmp/result/discovery.json ]; then mv /tmp/result/discovery.json /tmp/result/discovery_{template_name}.json; fi" + rename_cmd = f"\nif [ -f /tmp/result/discovery.json ]; then mv /tmp/result/discovery.json /tmp/result/discovery_{sanitized_template_name}.json; fi" if container.get("args"): container["args"] = [original_command + rename_cmd] diff --git a/coldpress_common/model.py b/coldpress_common/model.py index 807c4ed..f23bbb6 100644 --- a/coldpress_common/model.py +++ b/coldpress_common/model.py @@ -39,6 +39,7 @@ early rather than during job execution. """ +import re from typing import Any, Literal, Optional, Union from pydantic import ( BaseModel, @@ -484,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") From 6f5a76d771a41150dbc87094963342da12aec183 Mon Sep 17 00:00:00 2001 From: asanaullah Date: Wed, 29 Apr 2026 22:47:18 -0700 Subject: [PATCH 04/14] added workflow trigger for 0.2.1 --- .github/workflows/tests.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 515fec5..60bb640 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -2,9 +2,9 @@ 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: From 9115c1f07b06abc79898df59731dc236ccb19a47 Mon Sep 17 00:00:00 2001 From: asanaullah Date: Wed, 29 Apr 2026 22:53:16 -0700 Subject: [PATCH 05/14] added README --- README.md | 634 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 634 insertions(+) create mode 100644 README.md diff --git a/README.md b/README.md new file mode 100644 index 0000000..c21488a --- /dev/null +++ b/README.md @@ -0,0 +1,634 @@ +# Coldpress + +[![Tests](https://github.com/asanaullah/coldpress/workflows/Tests/badge.svg?branch=v0.2)](https://github.com/asanaullah/coldpress/actions/workflows/tests.yml) + +Coldpress is a prescriptive manifest generator that reduces the effort and expertise needed to deploy complex AI/HPC workloads on Kubernetes clusters. + +**Two-piece architecture:** + +1. **Admin** (`coldpress-setup`) - Generates cluster setup manifests (node labels, queues, namespaces, RBAC) +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 │ +└─────────────────────────────────────────────────────────────────┘ + ↓ +┌─────────────────────────────────────────────────────────────────┐ +│ Phase 1: Admin Setup (One-time) │ +├─────────────────────────────────────────────────────────────────┤ +│ 1. coldpress-setup generate cluster → manifests/cluster-*.yaml │ +│ → manifests/label-nodes-*.sh │ +│ │ +│ 2. ./manifests/label-nodes-*.sh (labels nodes for scheduling) │ +│ │ +│ 3. oc apply -f manifests/cluster-*.yaml │ +│ │ +│ 4. coldpress-setup generate project → manifests/project-*.yaml │ +│ │ +│ 5. oc apply -f manifests/project-*.yaml │ +│ │ +│ 6. coldpress-setup generate user → manifests/user-*.yaml │ +│ │ +│ 7. oc apply -f manifests/user-*.yaml │ +│ │ +│ Creates: Node labels, ClusterQueue, ResourceFlavors, │ +│ Namespaces, PVCs, RBAC │ +└─────────────────────────────────────────────────────────────────┘ + ↓ +┌─────────────────────────────────────────────────────────────────┐ +│ Phase 2: User Workflow (Repeatable) │ +├─────────────────────────────────────────────────────────────────┤ +│ 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 │ +│ │ +│ 3. User reviews generated manifest │ +│ │ +│ 4. ./run.sh applies manifest to cluster │ +│ │ +│ 5. Kueue schedules job when resources available │ +│ │ +│ 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 │ +│ │ +│ 7. ./logs.sh captures logs to PVC │ +│ │ +│ 8. ./explore.sh opens interactive shell to browse results │ +│ │ +│ 9. ./cp.sh copies results from PVC (optional) │ +│ │ +│ 10. ./cleanup.sh deletes resources (preserves results in PVC) │ +└─────────────────────────────────────────────────────────────────┘ +``` + +### What Coldpress Does + +**For Administrators (`coldpress-setup`):** +- Generates cluster-wide Kueue configuration (ClusterQueue, ResourceFlavors, LocalQueues) +- Generates project namespaces with PersistentVolumeClaims +- Generates user RBAC (RoleBindings) for job submission +- Generates node labeling scripts +- Outputs timestamped manifests for GitOps workflows + +**For Users (`coldpress`):** +- 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 +- Injects discovery init containers for hardware profiling +- Generates helper scripts for job lifecycle management +- Validates YAML schemas before generation + +## Installation + +Choose the installation method that fits your use case: + +| Use Case | Command | Activation Needed? | Best For | +|----------|---------|-------------------|----------| +| **Running jobs** | `./setup-env.sh --pipx` | ❌ No | End users, cluster users | +| **Running jobs + dev** | `./setup-env.sh --pipx-editable` | ❌ No | Users who also contribute | +| **Development** | `./setup-env.sh --uv` | ✅ Yes | Contributors, testing changes | + +### Quick Start + +```bash +# For end users (global install, no activation needed) +./setup-env.sh --pipx +coldpress --version + +# For developers (local venv, requires activation) +./setup-env.sh --uv +source .venv/bin/activate +coldpress --version +``` + +### For End Users (pipx) + +**pipx** installs Coldpress in an isolated environment with global CLI access - no activation needed. + +```bash +# Install with pipx (recommended for end users) +./setup-env.sh --pipx + +# Commands work globally, from any directory +coldpress --help +coldpress-setup --help +``` + +**Install in editable mode** (get updates as you pull from git): +```bash +./setup-env.sh --pipx-editable +``` + +**Manage installation:** +```bash +pipx upgrade coldpress # Upgrade to latest version +pipx reinstall coldpress # Reinstall +pipx uninstall coldpress # Remove completely +pipx list # Show installed packages +``` + +### For Developers (uv) + +**uv** provides fast, reproducible virtual environments for development work. + +```bash +# One-time setup (installs uv, creates venv, installs coldpress) +./setup-env.sh --uv +# or just: +./setup-env.sh # --uv is the default + +# Activate the environment +source .venv/bin/activate + +# Use coldpress (while venv is active) +coldpress --help + +# Or use without activation +.venv/bin/coldpress --help +``` + +**For subsequent sessions:** +```bash +source .venv/bin/activate +``` + +## Quick Start + +### For Administrators + +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 + ``` + +### For 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 + +```yaml +# Required +project: +output: +target: jobset | kubeflow | kuberay | kserve # Default: jobset + +# 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 +``` + +#### 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 + +```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" +``` + +### Macros + +Macros are placeholders in your job-spec.yaml that Coldpress automatically fills in when generating manifests. + +#### Available Macros + +**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` | + +**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}" +``` + +**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 +``` + +## Examples + +Coldpress supports multiple backend targets. Each example demonstrates a different use case and backend. + +### PyTorch DDP Training (JobSet) + +Distributed PyTorch training with 2 workers across 2 GPUs using JobSet. + +**Run:** +```bash +coldpress generate --intent examples/pytorch_ddp_training/intent_jobset.yaml +cd output/ddp-training-job/ +./run.sh +``` + +**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 +coldpress generate --intent examples/pytorch_ddp_training/intent_kubeflow.yaml +cd output/ddp-training-job/ +./run.sh +``` + +**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) + +### vLLM Inference (KServe) + +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) - 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) + +**Cluster tools:** +- `oc` CLI (tested with oc 4.17.0) + +## Environment Variables + +Customize directory locations with environment variables: + +**For `coldpress` (job generation):** +- `COLDPRESS_DISCOVERY_DIR` - Discovery templates directory (default: `discovery`) +- `COLDPRESS_PROJECT_DIR` - Project configs directory (default: `projects`) +- `COLDPRESS_OUTPUT_DIR` - Default output directory (default: `output`) + +**For `coldpress-setup` (manifest generation):** +- `COLDPRESS_MANIFESTS_DIR` - Manifest output directory (default: `manifests`) +- `COLDPRESS_CLUSTER_DIR` - Cluster configs directory (default: `cluster`) +- `COLDPRESS_USER_DIR` - User configs directory (default: `users`) + +**Example:** +```bash +export COLDPRESS_OUTPUT_DIR=jobs +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 +coldpress-setup generate project coldpress-project.yaml +# Outputs to: gitops/manifests/project-coldpress-project-*.yaml +``` + +## Resource Labels + +All generated resources 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.1` - Tracks Coldpress version +- `coldpress.io/job-id: {job_name}` - Job-specific identifier for compute resources + +**Query all Coldpress resources:** +```bash +oc get all -A -l app.kubernetes.io/managed-by=coldpress +``` + +**Delete resources by job:** +```bash +oc delete all -n namespace -l coldpress.io/job-id=job-name +``` + +## License + +See LICENSE file. From 02d07d31eb0a0ff2ba3c6885a82aa6a92aaa002f Mon Sep 17 00:00:00 2001 From: asanaullah Date: Wed, 29 Apr 2026 22:55:57 -0700 Subject: [PATCH 06/14] removed support for python 3.9 --- .github/workflows/tests.yml | 2 +- .gitignore | 2 +- README.md | 2 +- pyproject.toml | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 60bb640..739c99d 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -11,7 +11,7 @@ jobs: 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 diff --git a/.gitignore b/.gitignore index 99db0d9..c0d5885 100644 --- a/.gitignore +++ b/.gitignore @@ -84,5 +84,5 @@ users/test.yaml # Documentation docs/ -README.md +#README.md CHANGELOG.md diff --git a/README.md b/README.md index c21488a..d309598 100644 --- a/README.md +++ b/README.md @@ -581,7 +581,7 @@ coldpress/ - 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:** - `oc` CLI (tested with oc 4.17.0) diff --git a/pyproject.toml b/pyproject.toml index 86a4de3..d0629ff 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ name = "coldpress" 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", From 6d4a31122d40aa74de63c28d60931a97214db995 Mon Sep 17 00:00:00 2001 From: asanaullah Date: Wed, 29 Apr 2026 22:57:54 -0700 Subject: [PATCH 07/14] fixed outdated exit code test --- tests/test_exit_codes.sh | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) 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)" From 6d8caf81a826a5997e4d81cd6aef1c5bfdce1ae1 Mon Sep 17 00:00:00 2001 From: asanaullah Date: Wed, 29 Apr 2026 23:03:23 -0700 Subject: [PATCH 08/14] updated flowchart --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index d309598..45c63a6 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ Coldpress is a prescriptive manifest generator that reduces the effort and exper ``` ┌─────────────────────────────────────────────────────────────────┐ -│ Prerequisites: Kueue and JobSet operators must be installed │ +│ Prerequisites: Operators must be installed (see Requirements) │ └─────────────────────────────────────────────────────────────────┘ ↓ ┌─────────────────────────────────────────────────────────────────┐ From e5fed7ee8b387b4de4f51ad35fbf8cd66010501e Mon Sep 17 00:00:00 2001 From: asanaullah Date: Thu, 30 Apr 2026 10:42:45 -0700 Subject: [PATCH 09/14] cleanup --- CLAUDE.md | 115 ++++ .../ARG_OVERRIDE_README.md | 116 ---- .../pytorch_ddp_training/comp/job-spec.yaml | 75 --- .../pytorch_ddp_training/comp/jobset.yaml | 596 ------------------ .../pytorch_ddp_training/comp/pytorchjob.yaml | 131 ---- .../ARG_OVERRIDE_README.md | 105 --- examples/vllm_guidellm_benchmark/config.yaml | 10 - 7 files changed, 115 insertions(+), 1033 deletions(-) create mode 100644 CLAUDE.md delete mode 100644 examples/pytorch_ddp_training/ARG_OVERRIDE_README.md delete mode 100644 examples/pytorch_ddp_training/comp/job-spec.yaml delete mode 100644 examples/pytorch_ddp_training/comp/jobset.yaml delete mode 100644 examples/pytorch_ddp_training/comp/pytorchjob.yaml delete mode 100644 examples/pytorch_ray_training/ARG_OVERRIDE_README.md delete mode 100644 examples/vllm_guidellm_benchmark/config.yaml diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..61848d2 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,115 @@ +# Coldpress Project Documentation + +## What is Coldpress? + +Coldpress is a tool for AI/HPC workload orchestration on Kubernetes/OpenShift. It takes vanilla Kubernetes Job specifications and an intent file, then generates JobSet, Kubeflow (PyTorchJob, TFJob, MPIJob), or KServe (InferenceService) manifests with automatic configuration, discovery job injection, and result handling. + +## Architecture + +- **Input**: Vanilla K8s Job spec (job-spec.yaml) + Intent file (intent.yaml) +- **Processing**: Transform and inject configuration based on project settings +- **Output**: Backend-specific manifests (JobSet/Kubeflow/KServe) + helper scripts + discovery job injection +- **Features**: + - Discovery job injection to capture system state + - Automatic result handling and storage + - Prescriptive manifest generation with templating + - Helper scripts (run.sh, monitor.sh, explore.sh, cleanup.sh) + +## Critical Value Proposition Questions + +**This project needs clear answers to the following questions before it can justify adoption:** + +### 1. Core Value Questions + +- **Who asked for this?** What pain point drove the creation of coldpress? +- **Can you name 3 people who would choose this over Kubeflow + Helm?** +- **What's the 10-second pitch for why someone should adopt coldpress?** +- **What problem does coldpress solve that existing tools (Kubeflow, Helm, Kustomize) don't?** + +### 2. Abstraction Overhead + +- **Why write a 76-line K8s Job + 20-line intent when you could write a PyTorchJob directly?** +- **What's the productivity gain that justifies the learning curve?** +- **Does the intent.yaml abstraction reduce complexity or just move it?** +- **How does coldpress reduce time-to-first-job compared to standard Kubeflow?** + +### 3. Competing with Proven Tools + +- **Kubeflow already handles distributed training** - what does coldpress add? +- **Helm/Kustomize handle templating and configuration** - why a new format? +- **Standard K8s workflows are well-documented** - why create custom formats? +- **Lock-in risk**: What happens when users outgrow coldpress or need features it doesn't support? + +### 4. Target User Clarity + +- **ML researchers want simpler, not different YAML** - does coldpress simplify or complicate? +- **Platform teams want standard tools, not custom abstractions** - what's the platform team value? +- **Who is the primary user?** Data scientists? ML engineers? Platform engineers? Research teams? +- **What's their skill level assumption?** K8s experts? K8s beginners? ML practitioners with no K8s experience? + +### 5. Discovery Job Injection + +- **What problem does discovery job injection solve?** +- **Why capture system state for each job?** What's the use case? +- **How is this better than standard logging/monitoring tools?** +- **What do users do with the captured state?** +- **Does this add meaningful value or just create data overhead?** +- **How does discovery integrate with existing observability stacks?** +- **Can users achieve the same thing with Kubernetes audit logs or pod lifecycle hooks?** + +### 6. Automatic Result Handling + +- **What does "automatic result handling" mean specifically?** +- **How is this better than standard PVC + user-defined output paths?** +- **Does coldpress enforce a specific output structure?** Is that helpful or constraining? +- **How does this integrate with MLOps tools (MLflow, Weights & Biases, etc.)?** +- **What happens when users need custom result processing?** +- **Can users easily move results out of the coldpress ecosystem?** + +### 7. Prescriptive Manifest Generation + +- **What does "prescriptive" mean in this context?** +- **Who defines the prescriptions?** Platform team? Coldpress maintainers? Users? +- **How does prescriptive generation differ from templating (Helm/Kustomize)?** +- **Does "prescriptive" mean opinionated defaults?** What are those opinions based on? +- **Can users override prescriptions?** How much flexibility exists? +- **Does this reduce debugging surface or hide critical configuration?** +- **What's the escape hatch when prescriptions don't fit the use case?** + +### 8. Integration and Ecosystem + +- **How does coldpress fit into existing CI/CD pipelines?** +- **Can it coexist with existing Kubeflow/Argo workflows?** +- **What's the migration path from standard K8s/Kubeflow to coldpress?** +- **What's the migration path OUT of coldpress if needed?** + +### 9. Maintenance and Sustainability + +- **Who maintains coldpress long-term?** +- **How does it stay current with Kubernetes/Kubeflow API changes?** +- **What's the support model?** +- **Is this a research project or production-ready?** + +## Success Metrics + +Before investing further development, define: + +- **Adoption metric**: How many teams/users would validate success? +- **Time-to-value**: How much time should coldpress save per job vs. standard tools? +- **Complexity reduction**: Measurable reduction in YAML lines? Configuration errors? +- **User satisfaction**: What would make a user recommend coldpress? + +## Recommended Next Steps + +1. **Interview potential users** - validate the problem exists and coldpress solves it +2. **Competitive analysis** - document specific scenarios where coldpress wins vs. alternatives +3. **Create comparison demos** - side-by-side: coldpress vs. raw Kubeflow vs. Helm-templated +4. **Define the niche** - if coldpress isn't for everyone, who is it specifically for? +5. **Simplify or specialize** - either make it radically simpler OR focus on a specific use case where it's clearly superior + +## Notes for Contributors + +- All code changes should align with a clear answer to at least one of the value questions above +- Features should be justified by user research, not technical possibility +- Complexity should only be added when it measurably reduces user burden +- Standard K8s/Kubeflow patterns should be preferred unless coldpress provides clear improvement diff --git a/examples/pytorch_ddp_training/ARG_OVERRIDE_README.md b/examples/pytorch_ddp_training/ARG_OVERRIDE_README.md deleted file mode 100644 index 1a829c0..0000000 --- a/examples/pytorch_ddp_training/ARG_OVERRIDE_README.md +++ /dev/null @@ -1,116 +0,0 @@ -# Argument Override Pattern for Coldpress - -This example demonstrates the new argument override pattern for Coldpress intent files. - -## Problem - -The previous approach used environment variables with complex bash conditionals in the job-spec, making YAMLs unreadable and difficult to maintain. Single-node job-specs couldn't be directly run with `kubectl apply`. - -## Solution - -Use clean, directly-runnable job-specs with argument substitution from intent.yaml. - -## Key Features - -### 1. Clean Single-Node Job-Spec - -The `job-spec.yaml` is now a standard Kubernetes Job that can be directly applied: - -```yaml -command: - - python - - -m - - torch.distributed.run - - --nproc_per_node=2 - - train.py - - --dataset=mnist - - --epochs=50 - - --batch-size=128 -``` - -### 2. Intent-Based Argument Overrides - -Two patterns are supported in `intent.yaml`: - -#### Simple Replacement (existing args) -```yaml -args: - epochs: "100" # Replace existing --epochs value - batch-size: "256" # Replace existing --batch-size value -``` - -#### Positioned Insertion (new args) -```yaml -args: - nnodes: - value: "${REPLICAS}" - insert_after: nproc_per_node - - master_addr: - value: "${REPLICA_ddp-training_0}" - insert_after: node_rank -``` - -## Generated Output - -The JobSet generator produces: - -```yaml -command: - - python - - -m - - torch.distributed.run - - --nproc_per_node=2 - - --nnodes=2 # Inserted after nproc_per_node - - --node_rank=0 # Inserted after nnodes - - --master_addr=coldpress-... # Inserted after node_rank - - --master_port=29500 # Inserted after master_addr - - train.py - - --dataset=mnist - - --epochs=100 # Replaced from 50 - - --batch-size=256 # Replaced from 128 -``` - -## Macro Substitution - -The following macros are available: -- `${REPLICAS}` - Total replica count -- `${INDEX}` - Current replica index (0-based) -- `${REPLICA__N}` - DNS name of replica N -- `${NODE_ID}` - Assigned node ID (if using node pinning) - -## Implementation - -### Model Changes - -`coldpress_common/model.py`: -- Added `ArgOverride` Pydantic model -- Updated `TaskIntent.args` to support `Union[str, ArgOverride]` - -### Generator Changes - -Both `jobset_generator.py` and `kubeflow_generator.py`: -- Added `parse_arg()` function to parse CLI flags -- Enhanced `replace_or_add_arg()` with `insert_after`/`insert_before` support -- Updated arg processing to handle both `command` and `args` fields -- Handle both string values and ArgOverride objects - -## Example Files - -- `job-spec.yaml` - Clean single-node Job specification -- `intent_jobset.yaml` - JobSet intent with arg overrides -- `intent_kubeflow.yaml` - PyTorchJob intent (simpler, uses auto-injection) - -## Testing - -Generate manifests: -```bash -coldpress generate --intent intent_jobset.yaml -coldpress generate --intent intent_kubeflow.yaml -``` - -Verify args in output: -```bash -grep -A 20 "command:" output/ddp-training-job/jobset.yaml -grep -A 20 "command:" output/ddp-training-job/pytorchjob.yaml -``` diff --git a/examples/pytorch_ddp_training/comp/job-spec.yaml b/examples/pytorch_ddp_training/comp/job-spec.yaml deleted file mode 100644 index 7d3229c..0000000 --- a/examples/pytorch_ddp_training/comp/job-spec.yaml +++ /dev/null @@ -1,75 +0,0 @@ -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: ["bash", "-c"] - args: - - | - # Build torch.distributed.run command from env vars - CMD="python -m torch.distributed.run" - - # Add distributed params if env vars are set - [ -n "$NPROC_PER_NODE" ] && CMD="$CMD --nproc_per_node=$NPROC_PER_NODE" - [ -n "$NNODES" ] && CMD="$CMD --nnodes=$NNODES" - [ -n "$RANK" ] && CMD="$CMD --node_rank=$RANK" - [ -n "$MASTER_ADDR" ] && CMD="$CMD --master_addr=$MASTER_ADDR" - [ -n "$MASTER_PORT" ] && CMD="$CMD --master_port=$MASTER_PORT" - - # Training script and args - CMD="$CMD train.py" - CMD="$CMD --dataset=mnist" - CMD="$CMD --train-test-split=0.8" - CMD="$CMD --epochs=50" - CMD="$CMD --batch-size=128" - CMD="$CMD --hidden-size=4096" - CMD="$CMD --lr=0.01" - CMD="$CMD --output-dir=/results/checkpoints" - - echo "Running: $CMD" - exec $CMD - resources: - requests: - nvidia.com/gpu: "2" - memory: "16Gi" - cpu: "8" - limits: - nvidia.com/gpu: "2" - memory: "16Gi" - env: - - name: NCCL_DEBUG - value: "INFO" - - name: NPROC_PER_NODE - value: "2" - 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/comp/jobset.yaml b/examples/pytorch_ddp_training/comp/jobset.yaml deleted file mode 100644 index ce6702a..0000000 --- a/examples/pytorch_ddp_training/comp/jobset.yaml +++ /dev/null @@ -1,596 +0,0 @@ -apiVersion: jobset.x-k8s.io/v1alpha2 -kind: JobSet -metadata: - name: coldpress-ddp-training - namespace: coldpress-project - labels: - app.kubernetes.io/managed-by: coldpress - app.kubernetes.io/version: 0.2.0 - kueue.x-k8s.io/queue-name: coldpress-local-queue-coldpress-project - coldpress.io/job-id: coldpress-ddp-training - annotations: - coldpress.io/base-dir: coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821 - coldpress.io/storage-pvc: coldpress-project-storage -spec: - suspend: true - replicatedJobs: - - 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 - args: - - mkdir -p /data/coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821 - /data/coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821/task-0 - /data/coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821/task-1 - && echo Created directory /data/coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821 - with 2 task subdirectories - volumeMounts: - - name: storage - mountPath: /data - volumes: - - name: storage - persistentVolumeClaim: - claimName: coldpress-project-storage - - name: task-0 - replicas: 1 - template: - spec: - parallelism: 1 - completions: 1 - backoffLimit: 0 - template: - metadata: - labels: - app: task-0 - coldpress/gid: coldpress-ddp-training - spec: - restartPolicy: Never - containers: - - name: training - image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime - workingDir: /workspace - command: - - bash - - -c - args: - - '# Build torch.distributed.run command from env vars - - CMD="python -m torch.distributed.run" - - - # Add distributed params if env vars are set - - [ -n "$NPROC_PER_NODE" ] && CMD="$CMD --nproc_per_node=$NPROC_PER_NODE" - - [ -n "$NNODES" ] && CMD="$CMD --nnodes=$NNODES" - - [ -n "$RANK" ] && CMD="$CMD --node_rank=$RANK" - - [ -n "$MASTER_ADDR" ] && CMD="$CMD --master_addr=$MASTER_ADDR" - - [ -n "$MASTER_PORT" ] && CMD="$CMD --master_port=$MASTER_PORT" - - - # Training script and args - - CMD="$CMD train.py" - - CMD="$CMD --dataset=mnist" - - CMD="$CMD --train-test-split=0.8" - - CMD="$CMD --epochs=50" - - CMD="$CMD --batch-size=128" - - CMD="$CMD --hidden-size=4096" - - CMD="$CMD --lr=0.01" - - CMD="$CMD --output-dir=/results/checkpoints" - - - echo "Running: $CMD" - - exec $CMD - - ' - resources: - requests: - nvidia.com/gpu: '2' - memory: 16Gi - cpu: '8' - limits: - nvidia.com/gpu: '2' - memory: 16Gi - env: - - name: NCCL_DEBUG - value: INFO - - name: NPROC_PER_NODE - value: '2' - - name: NNODES - value: '2' - - name: MASTER_ADDR - value: coldpress-ddp-training-task-0-0-0.coldpress-ddp-training.coldpress-project.svc.cluster.local - - name: MASTER_PORT - value: '29500' - - name: RANK - value: '0' - - name: WORLD_SIZE - value: '2' - volumeMounts: - - name: results - mountPath: /results - subPath: coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821/task-0 - - name: dshm - mountPath: /dev/shm - - name: training-script - mountPath: /workspace/train.py - subPath: train.py - tolerations: &id001 - - operator: Exists - volumes: - - name: results - persistentVolumeClaim: - claimName: coldpress-project-storage - - name: dshm - emptyDir: - medium: Memory - sizeLimit: 16Gi - - name: training-script - configMap: - name: coldpress-ddp-training-files - initContainers: - - name: discovery - image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime - command: - - /bin/bash - - -c - args: - - "echo \"User System Snapshot - Capturing performance baseline with\ - \ dynamic benchmarks...\"\n\n# Create comprehensive Python script\ - \ in /tmp\ncat << 'EOF_SNAPSHOT' > /tmp/generate_snapshot.py\nimport\ - \ subprocess\nimport json\nimport time\nimport os\nfrom datetime import\ - \ datetime\n\ndef run_cmd(command):\n try:\n out = subprocess.run(command,\ - \ shell=True, capture_output=True, text=True, timeout=10)\n \ - \ return out.stdout.strip() if out.returncode == 0 else \"Not available\"\ - \n except Exception as e:\n return \"Not available\"\n\n\ - def cpu_benchmark():\n \"\"\"Simple CPU compute benchmark - measures\ - \ MFLOPS\"\"\"\n start = time.time()\n result = 0.0\n iterations\ - \ = 10000000\n for i in range(iterations):\n result += i\ - \ * 1.5 - 0.7\n elapsed = time.time() - start\n mflops = (iterations\ - \ * 2) / (elapsed * 1000000) # 2 ops per iteration\n return {\n\ - \ \"iterations\": iterations,\n \"time_seconds\": f\"\ - {elapsed:.4f}\",\n \"mflops\": f\"{mflops:.2f}\"\n }\n\n\ - def memory_bandwidth_benchmark():\n \"\"\"Memory sequential read/write\ - \ benchmark\"\"\"\n size_mb = 100\n size_bytes = size_mb * 1024\ - \ * 1024\n data = bytearray(size_bytes)\n\n # Write test\n \ - \ start = time.time()\n for i in range(0, size_bytes, 8):\n \ - \ data[i:i+8] = b'\\x00\\x01\\x02\\x03\\x04\\x05\\x06\\x07'\n\ - \ write_time = time.time() - start\n write_bw = (size_mb / write_time)\ - \ if write_time > 0 else 0\n\n # Read test\n start = time.time()\n\ - \ checksum = 0\n for i in range(0, size_bytes, 8):\n \ - \ checksum += data[i]\n read_time = time.time() - start\n read_bw\ - \ = (size_mb / read_time) if read_time > 0 else 0\n\n return {\n\ - \ \"test_size_mb\": size_mb,\n \"write_bandwidth_mbps\"\ - : f\"{write_bw:.2f}\",\n \"read_bandwidth_mbps\": f\"{read_bw:.2f}\"\ - ,\n \"write_time_sec\": f\"{write_time:.4f}\",\n \"\ - read_time_sec\": f\"{read_time:.4f}\"\n }\n\ndef disk_io_benchmark():\n\ - \ \"\"\"Disk I/O benchmark using /tmp\"\"\"\n test_file = \"\ - /tmp/io_bench_test.dat\"\n size_mb = 50\n block_size = 1024\ - \ * 1024 # 1MB blocks\n\n try:\n # Write test\n \ - \ start = time.time()\n with open(test_file, 'wb') as f:\n\ - \ for _ in range(size_mb):\n f.write(os.urandom(block_size))\n\ - \ f.flush()\n os.fsync(f.fileno())\n \ - \ write_time = time.time() - start\n write_bw = (size_mb /\ - \ write_time) if write_time > 0 else 0\n\n # Read test\n \ - \ start = time.time()\n with open(test_file, 'rb') as\ - \ f:\n while f.read(block_size):\n pass\n\ - \ read_time = time.time() - start\n read_bw = (size_mb\ - \ / read_time) if read_time > 0 else 0\n\n # Cleanup\n \ - \ os.remove(test_file)\n\n return {\n \"test_size_mb\"\ - : size_mb,\n \"write_bandwidth_mbps\": f\"{write_bw:.2f}\"\ - ,\n \"read_bandwidth_mbps\": f\"{read_bw:.2f}\",\n \ - \ \"write_time_sec\": f\"{write_time:.4f}\",\n \"\ - read_time_sec\": f\"{read_time:.4f}\"\n }\n except Exception\ - \ as e:\n return {\"error\": str(e)}\n\ndef context_switch_benchmark():\n\ - \ \"\"\"Measure context switch overhead using subprocess spawning\"\ - \"\"\n iterations = 100\n start = time.time()\n for _ in\ - \ range(iterations):\n subprocess.run(['true'], capture_output=True)\n\ - \ elapsed = time.time() - start\n avg_per_switch = (elapsed\ - \ / iterations) * 1000 # milliseconds\n\n return {\n \"\ - iterations\": iterations,\n \"total_time_sec\": f\"{elapsed:.4f}\"\ - ,\n \"avg_per_spawn_ms\": f\"{avg_per_switch:.4f}\"\n }\n\ - \nprint(\"Collecting static system information...\")\n\nsnapshot =\ - \ {\n \"timestamp\": datetime.now().isoformat(),\n \"system\"\ - : {\n \"cpu_model\": run_cmd(\"grep 'model name' /proc/cpuinfo\ - \ | head -1 | cut -d: -f2 | xargs\"),\n \"cpu_count\": run_cmd(\"\ - grep -c ^processor /proc/cpuinfo\"),\n \"cpu_cores_per_socket\"\ - : run_cmd(\"lscpu | grep 'Core(s) per socket' | awk '{print $NF}'\ - \ 2>/dev/null || echo 'N/A'\"),\n \"cpu_threads_per_core\"\ - : run_cmd(\"lscpu | grep 'Thread(s) per core' | awk '{print $NF}'\ - \ 2>/dev/null || echo 'N/A'\"),\n \"cpu_architecture\": run_cmd(\"\ - uname -m\"),\n \"cpu_current_freq_mhz\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_cur_freq\ - \ 2>/dev/null | awk '{print $1/1000}' || echo 'N/A'\"),\n \"\ - cpu_max_freq_mhz\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_max_freq\ - \ 2>/dev/null | awk '{print $1/1000}' || echo 'N/A'\"),\n \"\ - cpu_min_freq_mhz\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_min_freq\ - \ 2>/dev/null | awk '{print $1/1000}' || echo 'N/A'\"),\n \"\ - cache_l1d\": run_cmd(\"lscpu | grep 'L1d cache' | awk '{print $NF}'\ - \ 2>/dev/null || echo 'N/A'\"),\n \"cache_l1i\": run_cmd(\"\ - lscpu | grep 'L1i cache' | awk '{print $NF}' 2>/dev/null || echo 'N/A'\"\ - ),\n \"cache_l2\": run_cmd(\"lscpu | grep 'L2 cache' | awk\ - \ '{print $NF}' 2>/dev/null || echo 'N/A'\"),\n \"cache_l3\"\ - : run_cmd(\"lscpu | grep 'L3 cache' | awk '{print $NF}' 2>/dev/null\ - \ || echo 'N/A'\"),\n \"memory_total_gb\": run_cmd(\"grep MemTotal\ - \ /proc/meminfo | awk '{printf \\\"%.2f\\\", $2/1024/1024}'\"),\n\ - \ \"memory_available_gb\": run_cmd(\"grep MemAvailable /proc/meminfo\ - \ | awk '{printf \\\"%.2f\\\", $2/1024/1024}'\"),\n \"memory_free_gb\"\ - : run_cmd(\"grep MemFree /proc/meminfo | awk '{printf \\\"%.2f\\\"\ - , $2/1024/1024}'\"),\n \"numa_node_count\": run_cmd(\"ls -d\ - \ /sys/devices/system/node/node* 2>/dev/null | wc -l\"),\n \ - \ \"kernel_version\": run_cmd(\"uname -r\"),\n \"uptime_seconds\"\ - : run_cmd(\"cat /proc/uptime | awk '{print $1}'\")\n },\n \"\ - gpu\": {\n \"gpu_count\": run_cmd(\"nvidia-smi --query-gpu=name\ - \ --format=csv,noheader 2>/dev/null | wc -l\"),\n \"gpu_models\"\ - : run_cmd(\"nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null\"\ - ),\n \"gpu_memory_per_device_mb\": run_cmd(\"nvidia-smi --query-gpu=memory.total\ - \ --format=csv,noheader,nounits 2>/dev/null\"),\n \"gpu_driver_version\"\ - : run_cmd(\"nvidia-smi --query-gpu=driver_version --format=csv,noheader\ - \ 2>/dev/null | head -1\"),\n \"cuda_version\": run_cmd(\"\ - nvidia-smi | grep 'CUDA Version' | awk '{print $9}' 2>/dev/null ||\ - \ echo 'N/A'\")\n },\n \"network\": {\n \"interface_names\"\ - : run_cmd(\"ls /sys/class/net 2>/dev/null | grep -v lo | tr '\\\\\ - n' ', ' | sed 's/,$//'\"),\n \"rdma_devices\": run_cmd(\"ls\ - \ /sys/class/infiniband 2>/dev/null | tr '\\\\n' ', ' | sed 's/,$//'\ - \ || echo 'None'\"),\n \"rdma_device_count\": run_cmd(\"ls\ - \ /sys/class/infiniband 2>/dev/null | wc -l\"),\n \"primary_mtu\"\ - : run_cmd(\"cat /sys/class/net/eth0/mtu 2>/dev/null || echo 'N/A'\"\ - ),\n \"tcp_rmem\": run_cmd(\"cat /proc/sys/net/ipv4/tcp_rmem\ - \ 2>/dev/null || echo 'N/A'\"),\n \"tcp_wmem\": run_cmd(\"\ - cat /proc/sys/net/ipv4/tcp_wmem 2>/dev/null || echo 'N/A'\"),\n \ - \ \"tcp_congestion_control\": run_cmd(\"cat /proc/sys/net/ipv4/tcp_congestion_control\ - \ 2>/dev/null || echo 'N/A'\")\n },\n \"storage\": {\n \ - \ \"block_devices\": run_cmd(\"lsblk -nd -o NAME,SIZE 2>/dev/null\ - \ | grep -v loop\"),\n \"nvme_count\": run_cmd(\"ls /dev/nvme*\ - \ 2>/dev/null | grep -c nvme0n || echo '0'\"),\n \"tmpfs_size\"\ - : run_cmd(\"df -h /tmp 2>/dev/null | tail -1 | awk '{print $2}'\"\ - ),\n \"tmpfs_avail\": run_cmd(\"df -h /tmp 2>/dev/null | tail\ - \ -1 | awk '{print $4}'\"),\n \"io_scheduler\": run_cmd(\"\ - cat /sys/block/sda/queue/scheduler 2>/dev/null | grep -o '\\\\[.*\\\ - \\]' | tr -d '[]' || echo 'N/A'\")\n },\n \"performance_config\"\ - : {\n \"cpu_governor\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_governor\ - \ 2>/dev/null || echo 'N/A'\"),\n \"numa_balancing\": run_cmd(\"\ - cat /proc/sys/kernel/numa_balancing 2>/dev/null || echo 'N/A'\"),\n\ - \ \"transparent_hugepages\": run_cmd(\"cat /sys/kernel/mm/transparent_hugepage/enabled\ - \ 2>/dev/null | grep -o '\\\\[.*\\\\]' | tr -d '[]' || echo 'N/A'\"\ - ),\n \"swappiness\": run_cmd(\"cat /proc/sys/vm/swappiness\ - \ 2>/dev/null || echo 'N/A'\"),\n \"dirty_ratio\": run_cmd(\"\ - cat /proc/sys/vm/dirty_ratio 2>/dev/null || echo 'N/A'\"),\n \ - \ \"dirty_background_ratio\": run_cmd(\"cat /proc/sys/vm/dirty_background_ratio\ - \ 2>/dev/null || echo 'N/A'\"),\n \"max_map_count\": run_cmd(\"\ - cat /proc/sys/vm/max_map_count 2>/dev/null || echo 'N/A'\"),\n \ - \ \"file_max\": run_cmd(\"cat /proc/sys/fs/file-max 2>/dev/null\ - \ || echo 'N/A'\")\n },\n \"process_limits\": {\n \"\ - max_user_processes\": run_cmd(\"ulimit -u 2>/dev/null || echo 'N/A'\"\ - ),\n \"max_open_files\": run_cmd(\"ulimit -n 2>/dev/null ||\ - \ echo 'N/A'\"),\n \"max_stack_size_kb\": run_cmd(\"ulimit\ - \ -s 2>/dev/null || echo 'N/A'\"),\n \"max_memory_size_kb\"\ - : run_cmd(\"ulimit -m 2>/dev/null || echo 'N/A'\"),\n \"cgroup_cpu_quota\"\ - : run_cmd(\"cat /sys/fs/cgroup/cpu/cpu.cfs_quota_us 2>/dev/null ||\ - \ echo 'N/A'\"),\n \"cgroup_cpu_period\": run_cmd(\"cat /sys/fs/cgroup/cpu/cpu.cfs_period_us\ - \ 2>/dev/null || echo 'N/A'\"),\n \"cgroup_memory_limit_bytes\"\ - : run_cmd(\"cat /sys/fs/cgroup/memory/memory.limit_in_bytes 2>/dev/null\ - \ || echo 'N/A'\")\n }\n}\n\n# Dynamic Benchmarks\nprint(\"Running\ - \ CPU compute benchmark...\")\nsnapshot[\"benchmarks\"] = {\n \"\ - cpu_compute\": cpu_benchmark()\n}\n\nprint(\"Running memory bandwidth\ - \ benchmark...\")\nsnapshot[\"benchmarks\"][\"memory_bandwidth\"]\ - \ = memory_bandwidth_benchmark()\n\nprint(\"Running disk I/O benchmark...\"\ - )\nsnapshot[\"benchmarks\"][\"disk_io\"] = disk_io_benchmark()\n\n\ - print(\"Running context switch benchmark...\")\nsnapshot[\"benchmarks\"\ - ][\"context_switch\"] = context_switch_benchmark()\n\n# Write JSON\n\ - with open(\"/tmp/result/discovery.json\", \"w\") as f:\n json.dump(snapshot,\ - \ f, indent=2)\n\nprint(\"Snapshot generation complete!\")\nEOF_SNAPSHOT\n\ - \n# Execute snapshot generation\nmkdir -p /tmp/result\npython3 /tmp/generate_snapshot.py\n\ - \n# Clean up intermediate scripts\nrm -f /tmp/generate_snapshot.py\n\ - \necho \"User system snapshot complete. Results in discovery.json\"\ - \n\nif [ -f /tmp/result/discovery.json ]; then mv /tmp/result/discovery.json\ - \ /tmp/result/discovery_user_snapshot.json; fi" - volumeMounts: - - name: results - mountPath: /tmp/result - subPath: coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821/task-0 - dependsOn: - - name: mkdir - status: Complete - - name: task-1 - replicas: 1 - template: - spec: - parallelism: 1 - completions: 1 - backoffLimit: 0 - template: - metadata: - labels: - app: task-1 - coldpress/gid: coldpress-ddp-training - spec: - restartPolicy: Never - containers: - - name: training - image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime - workingDir: /workspace - command: - - bash - - -c - args: - - '# Build torch.distributed.run command from env vars - - CMD="python -m torch.distributed.run" - - - # Add distributed params if env vars are set - - [ -n "$NPROC_PER_NODE" ] && CMD="$CMD --nproc_per_node=$NPROC_PER_NODE" - - [ -n "$NNODES" ] && CMD="$CMD --nnodes=$NNODES" - - [ -n "$RANK" ] && CMD="$CMD --node_rank=$RANK" - - [ -n "$MASTER_ADDR" ] && CMD="$CMD --master_addr=$MASTER_ADDR" - - [ -n "$MASTER_PORT" ] && CMD="$CMD --master_port=$MASTER_PORT" - - - # Training script and args - - CMD="$CMD train.py" - - CMD="$CMD --dataset=mnist" - - CMD="$CMD --train-test-split=0.8" - - CMD="$CMD --epochs=50" - - CMD="$CMD --batch-size=128" - - CMD="$CMD --hidden-size=4096" - - CMD="$CMD --lr=0.01" - - CMD="$CMD --output-dir=/results/checkpoints" - - - echo "Running: $CMD" - - exec $CMD - - ' - resources: - requests: - nvidia.com/gpu: '2' - memory: 16Gi - cpu: '8' - limits: - nvidia.com/gpu: '2' - memory: 16Gi - env: - - name: NCCL_DEBUG - value: INFO - - name: NPROC_PER_NODE - value: '2' - - name: NNODES - value: '2' - - name: MASTER_ADDR - value: coldpress-ddp-training-task-0-0-0.coldpress-ddp-training.coldpress-project.svc.cluster.local - - name: MASTER_PORT - value: '29500' - - name: RANK - value: '1' - - name: WORLD_SIZE - value: '2' - volumeMounts: - - name: results - mountPath: /results - subPath: coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821/task-1 - - name: dshm - mountPath: /dev/shm - - name: training-script - mountPath: /workspace/train.py - subPath: train.py - tolerations: *id001 - volumes: - - name: results - persistentVolumeClaim: - claimName: coldpress-project-storage - - name: dshm - emptyDir: - medium: Memory - sizeLimit: 16Gi - - name: training-script - configMap: - name: coldpress-ddp-training-files - initContainers: - - name: discovery - image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime - command: - - /bin/bash - - -c - args: - - "echo \"User System Snapshot - Capturing performance baseline with\ - \ dynamic benchmarks...\"\n\n# Create comprehensive Python script\ - \ in /tmp\ncat << 'EOF_SNAPSHOT' > /tmp/generate_snapshot.py\nimport\ - \ subprocess\nimport json\nimport time\nimport os\nfrom datetime import\ - \ datetime\n\ndef run_cmd(command):\n try:\n out = subprocess.run(command,\ - \ shell=True, capture_output=True, text=True, timeout=10)\n \ - \ return out.stdout.strip() if out.returncode == 0 else \"Not available\"\ - \n except Exception as e:\n return \"Not available\"\n\n\ - def cpu_benchmark():\n \"\"\"Simple CPU compute benchmark - measures\ - \ MFLOPS\"\"\"\n start = time.time()\n result = 0.0\n iterations\ - \ = 10000000\n for i in range(iterations):\n result += i\ - \ * 1.5 - 0.7\n elapsed = time.time() - start\n mflops = (iterations\ - \ * 2) / (elapsed * 1000000) # 2 ops per iteration\n return {\n\ - \ \"iterations\": iterations,\n \"time_seconds\": f\"\ - {elapsed:.4f}\",\n \"mflops\": f\"{mflops:.2f}\"\n }\n\n\ - def memory_bandwidth_benchmark():\n \"\"\"Memory sequential read/write\ - \ benchmark\"\"\"\n size_mb = 100\n size_bytes = size_mb * 1024\ - \ * 1024\n data = bytearray(size_bytes)\n\n # Write test\n \ - \ start = time.time()\n for i in range(0, size_bytes, 8):\n \ - \ data[i:i+8] = b'\\x00\\x01\\x02\\x03\\x04\\x05\\x06\\x07'\n\ - \ write_time = time.time() - start\n write_bw = (size_mb / write_time)\ - \ if write_time > 0 else 0\n\n # Read test\n start = time.time()\n\ - \ checksum = 0\n for i in range(0, size_bytes, 8):\n \ - \ checksum += data[i]\n read_time = time.time() - start\n read_bw\ - \ = (size_mb / read_time) if read_time > 0 else 0\n\n return {\n\ - \ \"test_size_mb\": size_mb,\n \"write_bandwidth_mbps\"\ - : f\"{write_bw:.2f}\",\n \"read_bandwidth_mbps\": f\"{read_bw:.2f}\"\ - ,\n \"write_time_sec\": f\"{write_time:.4f}\",\n \"\ - read_time_sec\": f\"{read_time:.4f}\"\n }\n\ndef disk_io_benchmark():\n\ - \ \"\"\"Disk I/O benchmark using /tmp\"\"\"\n test_file = \"\ - /tmp/io_bench_test.dat\"\n size_mb = 50\n block_size = 1024\ - \ * 1024 # 1MB blocks\n\n try:\n # Write test\n \ - \ start = time.time()\n with open(test_file, 'wb') as f:\n\ - \ for _ in range(size_mb):\n f.write(os.urandom(block_size))\n\ - \ f.flush()\n os.fsync(f.fileno())\n \ - \ write_time = time.time() - start\n write_bw = (size_mb /\ - \ write_time) if write_time > 0 else 0\n\n # Read test\n \ - \ start = time.time()\n with open(test_file, 'rb') as\ - \ f:\n while f.read(block_size):\n pass\n\ - \ read_time = time.time() - start\n read_bw = (size_mb\ - \ / read_time) if read_time > 0 else 0\n\n # Cleanup\n \ - \ os.remove(test_file)\n\n return {\n \"test_size_mb\"\ - : size_mb,\n \"write_bandwidth_mbps\": f\"{write_bw:.2f}\"\ - ,\n \"read_bandwidth_mbps\": f\"{read_bw:.2f}\",\n \ - \ \"write_time_sec\": f\"{write_time:.4f}\",\n \"\ - read_time_sec\": f\"{read_time:.4f}\"\n }\n except Exception\ - \ as e:\n return {\"error\": str(e)}\n\ndef context_switch_benchmark():\n\ - \ \"\"\"Measure context switch overhead using subprocess spawning\"\ - \"\"\n iterations = 100\n start = time.time()\n for _ in\ - \ range(iterations):\n subprocess.run(['true'], capture_output=True)\n\ - \ elapsed = time.time() - start\n avg_per_switch = (elapsed\ - \ / iterations) * 1000 # milliseconds\n\n return {\n \"\ - iterations\": iterations,\n \"total_time_sec\": f\"{elapsed:.4f}\"\ - ,\n \"avg_per_spawn_ms\": f\"{avg_per_switch:.4f}\"\n }\n\ - \nprint(\"Collecting static system information...\")\n\nsnapshot =\ - \ {\n \"timestamp\": datetime.now().isoformat(),\n \"system\"\ - : {\n \"cpu_model\": run_cmd(\"grep 'model name' /proc/cpuinfo\ - \ | head -1 | cut -d: -f2 | xargs\"),\n \"cpu_count\": run_cmd(\"\ - grep -c ^processor /proc/cpuinfo\"),\n \"cpu_cores_per_socket\"\ - : run_cmd(\"lscpu | grep 'Core(s) per socket' | awk '{print $NF}'\ - \ 2>/dev/null || echo 'N/A'\"),\n \"cpu_threads_per_core\"\ - : run_cmd(\"lscpu | grep 'Thread(s) per core' | awk '{print $NF}'\ - \ 2>/dev/null || echo 'N/A'\"),\n \"cpu_architecture\": run_cmd(\"\ - uname -m\"),\n \"cpu_current_freq_mhz\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_cur_freq\ - \ 2>/dev/null | awk '{print $1/1000}' || echo 'N/A'\"),\n \"\ - cpu_max_freq_mhz\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_max_freq\ - \ 2>/dev/null | awk '{print $1/1000}' || echo 'N/A'\"),\n \"\ - cpu_min_freq_mhz\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_min_freq\ - \ 2>/dev/null | awk '{print $1/1000}' || echo 'N/A'\"),\n \"\ - cache_l1d\": run_cmd(\"lscpu | grep 'L1d cache' | awk '{print $NF}'\ - \ 2>/dev/null || echo 'N/A'\"),\n \"cache_l1i\": run_cmd(\"\ - lscpu | grep 'L1i cache' | awk '{print $NF}' 2>/dev/null || echo 'N/A'\"\ - ),\n \"cache_l2\": run_cmd(\"lscpu | grep 'L2 cache' | awk\ - \ '{print $NF}' 2>/dev/null || echo 'N/A'\"),\n \"cache_l3\"\ - : run_cmd(\"lscpu | grep 'L3 cache' | awk '{print $NF}' 2>/dev/null\ - \ || echo 'N/A'\"),\n \"memory_total_gb\": run_cmd(\"grep MemTotal\ - \ /proc/meminfo | awk '{printf \\\"%.2f\\\", $2/1024/1024}'\"),\n\ - \ \"memory_available_gb\": run_cmd(\"grep MemAvailable /proc/meminfo\ - \ | awk '{printf \\\"%.2f\\\", $2/1024/1024}'\"),\n \"memory_free_gb\"\ - : run_cmd(\"grep MemFree /proc/meminfo | awk '{printf \\\"%.2f\\\"\ - , $2/1024/1024}'\"),\n \"numa_node_count\": run_cmd(\"ls -d\ - \ /sys/devices/system/node/node* 2>/dev/null | wc -l\"),\n \ - \ \"kernel_version\": run_cmd(\"uname -r\"),\n \"uptime_seconds\"\ - : run_cmd(\"cat /proc/uptime | awk '{print $1}'\")\n },\n \"\ - gpu\": {\n \"gpu_count\": run_cmd(\"nvidia-smi --query-gpu=name\ - \ --format=csv,noheader 2>/dev/null | wc -l\"),\n \"gpu_models\"\ - : run_cmd(\"nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null\"\ - ),\n \"gpu_memory_per_device_mb\": run_cmd(\"nvidia-smi --query-gpu=memory.total\ - \ --format=csv,noheader,nounits 2>/dev/null\"),\n \"gpu_driver_version\"\ - : run_cmd(\"nvidia-smi --query-gpu=driver_version --format=csv,noheader\ - \ 2>/dev/null | head -1\"),\n \"cuda_version\": run_cmd(\"\ - nvidia-smi | grep 'CUDA Version' | awk '{print $9}' 2>/dev/null ||\ - \ echo 'N/A'\")\n },\n \"network\": {\n \"interface_names\"\ - : run_cmd(\"ls /sys/class/net 2>/dev/null | grep -v lo | tr '\\\\\ - n' ', ' | sed 's/,$//'\"),\n \"rdma_devices\": run_cmd(\"ls\ - \ /sys/class/infiniband 2>/dev/null | tr '\\\\n' ', ' | sed 's/,$//'\ - \ || echo 'None'\"),\n \"rdma_device_count\": run_cmd(\"ls\ - \ /sys/class/infiniband 2>/dev/null | wc -l\"),\n \"primary_mtu\"\ - : run_cmd(\"cat /sys/class/net/eth0/mtu 2>/dev/null || echo 'N/A'\"\ - ),\n \"tcp_rmem\": run_cmd(\"cat /proc/sys/net/ipv4/tcp_rmem\ - \ 2>/dev/null || echo 'N/A'\"),\n \"tcp_wmem\": run_cmd(\"\ - cat /proc/sys/net/ipv4/tcp_wmem 2>/dev/null || echo 'N/A'\"),\n \ - \ \"tcp_congestion_control\": run_cmd(\"cat /proc/sys/net/ipv4/tcp_congestion_control\ - \ 2>/dev/null || echo 'N/A'\")\n },\n \"storage\": {\n \ - \ \"block_devices\": run_cmd(\"lsblk -nd -o NAME,SIZE 2>/dev/null\ - \ | grep -v loop\"),\n \"nvme_count\": run_cmd(\"ls /dev/nvme*\ - \ 2>/dev/null | grep -c nvme0n || echo '0'\"),\n \"tmpfs_size\"\ - : run_cmd(\"df -h /tmp 2>/dev/null | tail -1 | awk '{print $2}'\"\ - ),\n \"tmpfs_avail\": run_cmd(\"df -h /tmp 2>/dev/null | tail\ - \ -1 | awk '{print $4}'\"),\n \"io_scheduler\": run_cmd(\"\ - cat /sys/block/sda/queue/scheduler 2>/dev/null | grep -o '\\\\[.*\\\ - \\]' | tr -d '[]' || echo 'N/A'\")\n },\n \"performance_config\"\ - : {\n \"cpu_governor\": run_cmd(\"cat /sys/devices/system/cpu/cpu0/cpufreq/scaling_governor\ - \ 2>/dev/null || echo 'N/A'\"),\n \"numa_balancing\": run_cmd(\"\ - cat /proc/sys/kernel/numa_balancing 2>/dev/null || echo 'N/A'\"),\n\ - \ \"transparent_hugepages\": run_cmd(\"cat /sys/kernel/mm/transparent_hugepage/enabled\ - \ 2>/dev/null | grep -o '\\\\[.*\\\\]' | tr -d '[]' || echo 'N/A'\"\ - ),\n \"swappiness\": run_cmd(\"cat /proc/sys/vm/swappiness\ - \ 2>/dev/null || echo 'N/A'\"),\n \"dirty_ratio\": run_cmd(\"\ - cat /proc/sys/vm/dirty_ratio 2>/dev/null || echo 'N/A'\"),\n \ - \ \"dirty_background_ratio\": run_cmd(\"cat /proc/sys/vm/dirty_background_ratio\ - \ 2>/dev/null || echo 'N/A'\"),\n \"max_map_count\": run_cmd(\"\ - cat /proc/sys/vm/max_map_count 2>/dev/null || echo 'N/A'\"),\n \ - \ \"file_max\": run_cmd(\"cat /proc/sys/fs/file-max 2>/dev/null\ - \ || echo 'N/A'\")\n },\n \"process_limits\": {\n \"\ - max_user_processes\": run_cmd(\"ulimit -u 2>/dev/null || echo 'N/A'\"\ - ),\n \"max_open_files\": run_cmd(\"ulimit -n 2>/dev/null ||\ - \ echo 'N/A'\"),\n \"max_stack_size_kb\": run_cmd(\"ulimit\ - \ -s 2>/dev/null || echo 'N/A'\"),\n \"max_memory_size_kb\"\ - : run_cmd(\"ulimit -m 2>/dev/null || echo 'N/A'\"),\n \"cgroup_cpu_quota\"\ - : run_cmd(\"cat /sys/fs/cgroup/cpu/cpu.cfs_quota_us 2>/dev/null ||\ - \ echo 'N/A'\"),\n \"cgroup_cpu_period\": run_cmd(\"cat /sys/fs/cgroup/cpu/cpu.cfs_period_us\ - \ 2>/dev/null || echo 'N/A'\"),\n \"cgroup_memory_limit_bytes\"\ - : run_cmd(\"cat /sys/fs/cgroup/memory/memory.limit_in_bytes 2>/dev/null\ - \ || echo 'N/A'\")\n }\n}\n\n# Dynamic Benchmarks\nprint(\"Running\ - \ CPU compute benchmark...\")\nsnapshot[\"benchmarks\"] = {\n \"\ - cpu_compute\": cpu_benchmark()\n}\n\nprint(\"Running memory bandwidth\ - \ benchmark...\")\nsnapshot[\"benchmarks\"][\"memory_bandwidth\"]\ - \ = memory_bandwidth_benchmark()\n\nprint(\"Running disk I/O benchmark...\"\ - )\nsnapshot[\"benchmarks\"][\"disk_io\"] = disk_io_benchmark()\n\n\ - print(\"Running context switch benchmark...\")\nsnapshot[\"benchmarks\"\ - ][\"context_switch\"] = context_switch_benchmark()\n\n# Write JSON\n\ - with open(\"/tmp/result/discovery.json\", \"w\") as f:\n json.dump(snapshot,\ - \ f, indent=2)\n\nprint(\"Snapshot generation complete!\")\nEOF_SNAPSHOT\n\ - \n# Execute snapshot generation\nmkdir -p /tmp/result\npython3 /tmp/generate_snapshot.py\n\ - \n# Clean up intermediate scripts\nrm -f /tmp/generate_snapshot.py\n\ - \necho \"User system snapshot complete. Results in discovery.json\"\ - \n\nif [ -f /tmp/result/discovery.json ]; then mv /tmp/result/discovery.json\ - \ /tmp/result/discovery_user_snapshot.json; fi" - volumeMounts: - - name: results - mountPath: /tmp/result - subPath: coldpress-project/coldpress_results/ddp-training-0006ac7f-20260415_235821/task-1 - dependsOn: - - name: mkdir - status: Complete - successPolicy: - operator: All - targetReplicatedJobs: - - task-0 - - task-1 diff --git a/examples/pytorch_ddp_training/comp/pytorchjob.yaml b/examples/pytorch_ddp_training/comp/pytorchjob.yaml deleted file mode 100644 index d8c5b22..0000000 --- a/examples/pytorch_ddp_training/comp/pytorchjob.yaml +++ /dev/null @@ -1,131 +0,0 @@ -apiVersion: kubeflow.org/v1 -kind: PyTorchJob -metadata: - name: coldpress-ddp-training - namespace: coldpress-project - labels: - app.kubernetes.io/managed-by: coldpress - app.kubernetes.io/version: 0.2.0 - coldpress.io/job-id: coldpress-ddp-training - kueue.x-k8s.io/queue-name: coldpress-local-queue-coldpress-project - annotations: - coldpress.io/base-dir: coldpress-project/coldpress_results/ddp-training-050895d3-20260416_000443 - coldpress.io/storage-pvc: coldpress-project-storage - kueue.x-k8s.io/wait-for-pods-ready-timeout: 20m -spec: - suspend: true - pytorchReplicaSpecs: - Master: - replicas: 1 - restartPolicy: OnFailure - template: &id001 - metadata: - annotations: - sidecar.istio.io/inject: 'false' - labels: - app.kubernetes.io/managed-by: coldpress - app.kubernetes.io/version: 0.2.0 - coldpress.io/job-id: coldpress-ddp-training - spec: - restartPolicy: OnFailure - containers: - - name: pytorch - image: pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime - workingDir: /workspace - command: - - bash - - -c - args: - - '# Build torch.distributed.run command from env vars - - CMD="python -m torch.distributed.run" - - - # Add distributed params if env vars are set - - [ -n "$NPROC_PER_NODE" ] && CMD="$CMD --nproc_per_node=$NPROC_PER_NODE" - - [ -n "$NNODES" ] && CMD="$CMD --nnodes=$NNODES" - - [ -n "$RANK" ] && CMD="$CMD --node_rank=$RANK" - - [ -n "$MASTER_ADDR" ] && CMD="$CMD --master_addr=$MASTER_ADDR" - - [ -n "$MASTER_PORT" ] && CMD="$CMD --master_port=$MASTER_PORT" - - - # Training script and args - - CMD="$CMD train.py" - - CMD="$CMD --dataset=mnist" - - CMD="$CMD --train-test-split=0.8" - - CMD="$CMD --epochs=50" - - CMD="$CMD --batch-size=128" - - CMD="$CMD --hidden-size=4096" - - CMD="$CMD --lr=0.01" - - CMD="$CMD --output-dir=/results/checkpoints" - - - echo "Running: $CMD" - - exec $CMD - - ' - - --nnodes=2 - resources: - requests: - nvidia.com/gpu: '2' - memory: 16Gi - cpu: '8' - limits: - nvidia.com/gpu: '2' - memory: 16Gi - env: - - name: NCCL_DEBUG - value: INFO - - name: NPROC_PER_NODE - value: '2' - volumeMounts: - - name: results - mountPath: /results - - name: dshm - mountPath: /dev/shm - - name: training-script - mountPath: /workspace/train.py - subPath: train.py - tolerations: - - operator: Exists - volumes: - - name: results - persistentVolumeClaim: - claimName: coldpress-project-storage - - name: dshm - emptyDir: - medium: Memory - sizeLimit: 16Gi - - name: training-script - configMap: - name: coldpress-ddp-training-files - initContainers: - - name: mkdir - image: registry.access.redhat.com/ubi9/ubi-minimal:latest - command: - - sh - - -c - args: - - mkdir -p /results/coldpress-project/coldpress_results/ddp-training-050895d3-20260416_000443 - && echo Created directory /results/coldpress-project/coldpress_results/ddp-training-050895d3-20260416_000443 - volumeMounts: - - name: results - mountPath: /results - Worker: - replicas: 1 - restartPolicy: OnFailure - template: *id001 diff --git a/examples/pytorch_ray_training/ARG_OVERRIDE_README.md b/examples/pytorch_ray_training/ARG_OVERRIDE_README.md deleted file mode 100644 index 884728b..0000000 --- a/examples/pytorch_ray_training/ARG_OVERRIDE_README.md +++ /dev/null @@ -1,105 +0,0 @@ -# Argument Override Pattern for Ray Training Example - -This example demonstrates the argument override pattern for Ray Train jobs using KubeRay, matching the approach used in the PyTorch DDP example. - -## Updates Made - -### 1. Clean Single-Node Job-Spec - -The `job-spec.yaml` now uses clean command/args (no bash conditionals): - -```yaml -command: - - python - - train.py - - --dataset=mnist - - --epochs=50 - - --batch-size=128 - - --hidden-size=4096 - - --lr=0.01 - - --num-workers=1 # Single-node default - - --gpus-per-worker=2 - - --cpus-per-worker=4 -``` - -**Before:** Used bash script with env var conditionals (complex, hard to read) -**After:** Direct command with arguments (clean, directly runnable) - -### 2. Intent-Based Argument Overrides - -The `intent_kuberay.yaml` now uses `args` instead of `env`: - -```yaml -tasks: - - name: ray-training - replicas: 2 # 1 head + 1 worker = 4 GPUs total - - args: - # Scale Ray Train workers to match pod count - num-workers: "${REPLICAS}" # Overrides from 1 to 2 - # gpus-per-worker and cpus-per-worker use defaults from job-spec -``` - -**Before:** Used `env` with environment variable overrides -**After:** Uses `args` with macro substitution (consistent with DDP/Kubeflow patterns) - -### 3. Generator Updates - -Updated `kuberay_generator.py` to: -- Import `parse_arg` and `replace_or_add_arg` from `jobset_generator` -- Apply arg overrides before building the RayJob entrypoint -- Handle both simple string values and `ArgOverride` objects with positioning -- Process args for both head and worker containers - -## Generated Output - -The RayJob manifest now has a clean entrypoint: - -```yaml -entrypoint: cd /results/.../workspace && 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=2 # Overridden from 1 - --gpus-per-worker=2 # Default preserved - --cpus-per-worker=4 # Default preserved -``` - -## Macro Substitution - -Available macros for RayJob: -- `${REPLICAS}` - Total replica count (head + workers) -- `${INDEX}` - Container index (0 for head, 1+ for workers) -- `${TASK_NAME}` - Task name from intent - -## Consistency with DDP Example - -Both examples now follow the same pattern: - -| Aspect | DDP Example | Ray Example | -|--------|-------------|-------------| -| job-spec | Clean command/args | Clean command/args | -| intent | Args with positioning | Args with macros | -| Generator | `jobset_generator.py` | `kuberay_generator.py` | -| Override function | `replace_or_add_arg()` | `replace_or_add_arg()` | - -## Testing - -Generate and verify: -```bash -coldpress generate --intent examples/pytorch_ray_training/intent_kuberay.yaml -grep "num-workers" output/ray-training-job/rayjob.yaml -# Should show: --num-workers=2 (overridden from 1) -``` - -## Benefits - -1. **Readable job-specs**: No bash scripts, directly runnable -2. **Consistent pattern**: Same approach across JobSet, Kubeflow, KubeRay -3. **Type-safe**: Pydantic validation of arg overrides -4. **Flexible**: Supports both simple replacement and positioned insertion -5. **Maintainable**: Clear separation between single-node defaults and multi-node scaling 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 From b47b4d562356a2ea5ca2b30a865bfe87882ac46b Mon Sep 17 00:00:00 2001 From: asanaullah Date: Tue, 5 May 2026 16:01:49 -0700 Subject: [PATCH 10/14] resolved conflicts with the coding guidelines --- coldpress/__init__.py | 2 +- coldpress/kubeflow_generator.py | 9 ++- coldpress/utils.py | 108 +++++++++++++++++++++++++------- coldpress_setup/__init__.py | 2 +- 4 files changed, 94 insertions(+), 27 deletions(-) diff --git a/coldpress/__init__.py b/coldpress/__init__.py index b59d2ff..5dea243 100644 --- a/coldpress/__init__.py +++ b/coldpress/__init__.py @@ -1,4 +1,4 @@ # Generated 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/kubeflow_generator.py b/coldpress/kubeflow_generator.py index 69235c5..54bf31b 100644 --- a/coldpress/kubeflow_generator.py +++ b/coldpress/kubeflow_generator.py @@ -230,6 +230,12 @@ def generate_pytorchjob_from_intent( 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", @@ -297,7 +303,8 @@ def generate_pytorchjob_from_intent( if storage_volume_name: discovery_init = build_discovery_init_container( - discovery_template_path, base_dir, storage_volume_name + discovery_template_path, base_dir, storage_volume_name, + task_id=None, per_pod_directory=True ) if discovery_init: init_containers.append(discovery_init) diff --git a/coldpress/utils.py b/coldpress/utils.py index 82e9f27..55d471b 100644 --- a/coldpress/utils.py +++ b/coldpress/utils.py @@ -413,6 +413,7 @@ def build_discovery_init_container( base_dir: str, storage_volume_name: str, task_id: int | None = None, + per_pod_directory: bool = False, ): """ Build discovery init container from template. @@ -422,6 +423,7 @@ def build_discovery_init_container( 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 @@ -451,32 +453,90 @@ def build_discovery_init_container( container = copy.deepcopy(containers[0]) container["name"] = "discovery" - # Construct result path - if task_id is not None: - # JobSet mode - task-specific path - result_path = f"{base_dir}/task-{task_id}" + 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 /tmp/result references with current directory + modified_command = original_command.replace("/tmp/result/", "./") + # 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: - # KubeRay/Kubeflow mode - base directory only - result_path = base_dir - - # Update volume mounts to use PVC - container["volumeMounts"] = [ - { - "name": storage_volume_name, - "mountPath": "/tmp/result", - "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 /tmp/result/discovery.json ]; then mv /tmp/result/discovery.json /tmp/result/discovery_{sanitized_template_name}.json; fi" + # 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": "/tmp/result", + "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 /tmp/result/discovery.json ]; then mv /tmp/result/discovery.json /tmp/result/discovery_{sanitized_template_name}.json; fi" - if container.get("args"): - container["args"] = [original_command + rename_cmd] + if container.get("args"): + container["args"] = [original_command + rename_cmd] return container except (FileNotFoundError, yaml.YAMLError, KeyError, IndexError) as e: diff --git a/coldpress_setup/__init__.py b/coldpress_setup/__init__.py index 0c535bd..7d51242 100644 --- a/coldpress_setup/__init__.py +++ b/coldpress_setup/__init__.py @@ -1,4 +1,4 @@ # Generated by: Claude Sonnet 4.5 """Coldpress Config - Prescriptive cluster configuration.""" -__version__ = "0.2.0" +__version__ = "0.2.1" From a11c81ad0987fa936bdeb73764a3d9cdcc12af8c Mon Sep 17 00:00:00 2001 From: asanaullah Date: Tue, 5 May 2026 16:11:48 -0700 Subject: [PATCH 11/14] Resolved Kubeflow manifest generation issues - Fixed version strings: updated coldpress/__init__.py and coldpress_setup/__init__.py from 0.2.0 to 0.2.1 - Fixed missing subPath in PyTorchJob main container results volume mount (kubeflow_generator.py) - Implemented per-pod discovery directories using downward API with KubeRay naming convention: * Master writes to base directory * Workers write to worker-{index}/ subdirectories - Added CLAUDE.md to .gitignore and removed from version control Co-Authored-By: Claude Sonnet 4.5 --- .gitignore | 1 + CLAUDE.md | 115 ----------------------------------------------------- 2 files changed, 1 insertion(+), 115 deletions(-) delete mode 100644 CLAUDE.md diff --git a/.gitignore b/.gitignore index c0d5885..fc238f7 100644 --- a/.gitignore +++ b/.gitignore @@ -86,3 +86,4 @@ users/test.yaml docs/ #README.md CHANGELOG.md +CLAUDE.md diff --git a/CLAUDE.md b/CLAUDE.md deleted file mode 100644 index 61848d2..0000000 --- a/CLAUDE.md +++ /dev/null @@ -1,115 +0,0 @@ -# Coldpress Project Documentation - -## What is Coldpress? - -Coldpress is a tool for AI/HPC workload orchestration on Kubernetes/OpenShift. It takes vanilla Kubernetes Job specifications and an intent file, then generates JobSet, Kubeflow (PyTorchJob, TFJob, MPIJob), or KServe (InferenceService) manifests with automatic configuration, discovery job injection, and result handling. - -## Architecture - -- **Input**: Vanilla K8s Job spec (job-spec.yaml) + Intent file (intent.yaml) -- **Processing**: Transform and inject configuration based on project settings -- **Output**: Backend-specific manifests (JobSet/Kubeflow/KServe) + helper scripts + discovery job injection -- **Features**: - - Discovery job injection to capture system state - - Automatic result handling and storage - - Prescriptive manifest generation with templating - - Helper scripts (run.sh, monitor.sh, explore.sh, cleanup.sh) - -## Critical Value Proposition Questions - -**This project needs clear answers to the following questions before it can justify adoption:** - -### 1. Core Value Questions - -- **Who asked for this?** What pain point drove the creation of coldpress? -- **Can you name 3 people who would choose this over Kubeflow + Helm?** -- **What's the 10-second pitch for why someone should adopt coldpress?** -- **What problem does coldpress solve that existing tools (Kubeflow, Helm, Kustomize) don't?** - -### 2. Abstraction Overhead - -- **Why write a 76-line K8s Job + 20-line intent when you could write a PyTorchJob directly?** -- **What's the productivity gain that justifies the learning curve?** -- **Does the intent.yaml abstraction reduce complexity or just move it?** -- **How does coldpress reduce time-to-first-job compared to standard Kubeflow?** - -### 3. Competing with Proven Tools - -- **Kubeflow already handles distributed training** - what does coldpress add? -- **Helm/Kustomize handle templating and configuration** - why a new format? -- **Standard K8s workflows are well-documented** - why create custom formats? -- **Lock-in risk**: What happens when users outgrow coldpress or need features it doesn't support? - -### 4. Target User Clarity - -- **ML researchers want simpler, not different YAML** - does coldpress simplify or complicate? -- **Platform teams want standard tools, not custom abstractions** - what's the platform team value? -- **Who is the primary user?** Data scientists? ML engineers? Platform engineers? Research teams? -- **What's their skill level assumption?** K8s experts? K8s beginners? ML practitioners with no K8s experience? - -### 5. Discovery Job Injection - -- **What problem does discovery job injection solve?** -- **Why capture system state for each job?** What's the use case? -- **How is this better than standard logging/monitoring tools?** -- **What do users do with the captured state?** -- **Does this add meaningful value or just create data overhead?** -- **How does discovery integrate with existing observability stacks?** -- **Can users achieve the same thing with Kubernetes audit logs or pod lifecycle hooks?** - -### 6. Automatic Result Handling - -- **What does "automatic result handling" mean specifically?** -- **How is this better than standard PVC + user-defined output paths?** -- **Does coldpress enforce a specific output structure?** Is that helpful or constraining? -- **How does this integrate with MLOps tools (MLflow, Weights & Biases, etc.)?** -- **What happens when users need custom result processing?** -- **Can users easily move results out of the coldpress ecosystem?** - -### 7. Prescriptive Manifest Generation - -- **What does "prescriptive" mean in this context?** -- **Who defines the prescriptions?** Platform team? Coldpress maintainers? Users? -- **How does prescriptive generation differ from templating (Helm/Kustomize)?** -- **Does "prescriptive" mean opinionated defaults?** What are those opinions based on? -- **Can users override prescriptions?** How much flexibility exists? -- **Does this reduce debugging surface or hide critical configuration?** -- **What's the escape hatch when prescriptions don't fit the use case?** - -### 8. Integration and Ecosystem - -- **How does coldpress fit into existing CI/CD pipelines?** -- **Can it coexist with existing Kubeflow/Argo workflows?** -- **What's the migration path from standard K8s/Kubeflow to coldpress?** -- **What's the migration path OUT of coldpress if needed?** - -### 9. Maintenance and Sustainability - -- **Who maintains coldpress long-term?** -- **How does it stay current with Kubernetes/Kubeflow API changes?** -- **What's the support model?** -- **Is this a research project or production-ready?** - -## Success Metrics - -Before investing further development, define: - -- **Adoption metric**: How many teams/users would validate success? -- **Time-to-value**: How much time should coldpress save per job vs. standard tools? -- **Complexity reduction**: Measurable reduction in YAML lines? Configuration errors? -- **User satisfaction**: What would make a user recommend coldpress? - -## Recommended Next Steps - -1. **Interview potential users** - validate the problem exists and coldpress solves it -2. **Competitive analysis** - document specific scenarios where coldpress wins vs. alternatives -3. **Create comparison demos** - side-by-side: coldpress vs. raw Kubeflow vs. Helm-templated -4. **Define the niche** - if coldpress isn't for everyone, who is it specifically for? -5. **Simplify or specialize** - either make it radically simpler OR focus on a specific use case where it's clearly superior - -## Notes for Contributors - -- All code changes should align with a clear answer to at least one of the value questions above -- Features should be justified by user research, not technical possibility -- Complexity should only be added when it measurably reduces user burden -- Standard K8s/Kubeflow patterns should be preferred unless coldpress provides clear improvement From 0b44ec707013d61c83eb1602ad2e01ecb53eb524 Mon Sep 17 00:00:00 2001 From: asanaullah Date: Tue, 5 May 2026 16:17:00 -0700 Subject: [PATCH 12/14] fixed format issues --- coldpress/kubeflow_generator.py | 7 ++++-- coldpress/utils.py | 38 +++++++++++++++++---------------- 2 files changed, 25 insertions(+), 20 deletions(-) diff --git a/coldpress/kubeflow_generator.py b/coldpress/kubeflow_generator.py index 54bf31b..06fe709 100644 --- a/coldpress/kubeflow_generator.py +++ b/coldpress/kubeflow_generator.py @@ -303,8 +303,11 @@ def generate_pytorchjob_from_intent( 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 + discovery_template_path, + base_dir, + storage_volume_name, + task_id=None, + per_pod_directory=True, ) if discovery_init: init_containers.append(discovery_init) diff --git a/coldpress/utils.py b/coldpress/utils.py index 55d471b..a63920f 100644 --- a/coldpress/utils.py +++ b/coldpress/utils.py @@ -467,24 +467,26 @@ def build_discovery_init_container( # 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']" - } - } - } - ]) + 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}/ From dcb5425b644d981609a74375d49c9b64a4f35ccb Mon Sep 17 00:00:00 2001 From: asanaullah Date: Tue, 5 May 2026 16:22:04 -0700 Subject: [PATCH 13/14] fixed false positive in security testing with /tmp paths inside pod --- coldpress/utils.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/coldpress/utils.py b/coldpress/utils.py index a63920f..17426c9 100644 --- a/coldpress/utils.py +++ b/coldpress/utils.py @@ -14,6 +14,9 @@ # 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: """ @@ -504,8 +507,10 @@ def build_discovery_init_container( fi mkdir -p $DISCOVERY_DIR && cd $DISCOVERY_DIR && """ - # Replace /tmp/result references with current directory - modified_command = original_command.replace("/tmp/result/", "./") + # 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" @@ -525,7 +530,7 @@ def build_discovery_init_container( container["volumeMounts"] = [ { "name": storage_volume_name, - "mountPath": "/tmp/result", + "mountPath": DISCOVERY_MOUNT_PATH, "subPath": result_path, } ] @@ -535,7 +540,7 @@ def build_discovery_init_container( 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_{sanitized_template_name}.json; fi" + 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] From 5b7c273bdd11bbaec9701746dcff300b9deb405b Mon Sep 17 00:00:00 2001 From: asanaullah Date: Tue, 5 May 2026 16:27:04 -0700 Subject: [PATCH 14/14] added/update headers to reflect Claude usage --- coldpress/__init__.py | 2 +- coldpress/cli.py | 2 +- coldpress/constants.py | 2 +- coldpress/jobset_generator.py | 2 +- coldpress/kubeflow_generator.py | 2 +- coldpress/kuberay_generator.py | 2 +- coldpress/script_gen.py | 2 +- coldpress/utils.py | 2 +- coldpress_setup/__init__.py | 2 +- coldpress_setup/cli.py | 2 +- coldpress_setup/generator.py | 2 +- 11 files changed, 11 insertions(+), 11 deletions(-) diff --git a/coldpress/__init__.py b/coldpress/__init__.py index 5dea243..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.1" diff --git a/coldpress/cli.py b/coldpress/cli.py index f8cbe20..a47fc26 100644 --- a/coldpress/cli.py +++ b/coldpress/cli.py @@ -1,4 +1,4 @@ -# Generated by: Claude Sonnet 4.5 +# Assisted by: Claude Sonnet 4.5 """Coldpress CLI - Generate JobSet and Kubeflow manifests for AI/HPC workloads. Configuration Priority: diff --git a/coldpress/constants.py b/coldpress/constants.py index 68533a2..39a8da1 100644 --- a/coldpress/constants.py +++ b/coldpress/constants.py @@ -1,4 +1,4 @@ -# Generated by: Claude Sonnet 4.5 +# Assisted by: Claude Sonnet 4.5 """Coldpress constants - centralized configuration values.""" # Version diff --git a/coldpress/jobset_generator.py b/coldpress/jobset_generator.py index 5ef9c2a..0018541 100644 --- a/coldpress/jobset_generator.py +++ b/coldpress/jobset_generator.py @@ -1,4 +1,4 @@ -# Generated by: Claude Sonnet 4.5 +# Assisted by: Claude Sonnet 4.5 """JobSet generation from vanilla k8s Jobs + intent.yaml. Generates JobSet manifests directly from: diff --git a/coldpress/kubeflow_generator.py b/coldpress/kubeflow_generator.py index 06fe709..6fe2f0a 100644 --- a/coldpress/kubeflow_generator.py +++ b/coldpress/kubeflow_generator.py @@ -1,4 +1,4 @@ -# Generated by: Claude Sonnet 4.5 +# Assisted by: Claude Sonnet 4.5 """Kubeflow manifest generation for Coldpress jobs.""" import copy diff --git a/coldpress/kuberay_generator.py b/coldpress/kuberay_generator.py index e699a3e..6377189 100644 --- a/coldpress/kuberay_generator.py +++ b/coldpress/kuberay_generator.py @@ -1,4 +1,4 @@ -# Generated by: Claude Sonnet 4.5 +# Assisted by: Claude Sonnet 4.5 """KubeRay manifest generation for Coldpress jobs. Transforms single-node Ray Train jobs into distributed RayJob manifests. diff --git a/coldpress/script_gen.py b/coldpress/script_gen.py index 5b65b83..a76c558 100644 --- a/coldpress/script_gen.py +++ b/coldpress/script_gen.py @@ -1,4 +1,4 @@ -# Generated by: Claude Sonnet 4.5 +# Assisted by: Claude Sonnet 4.5 """Bash script generation for Coldpress jobs.""" import os diff --git a/coldpress/utils.py b/coldpress/utils.py index 17426c9..2d80bf6 100644 --- a/coldpress/utils.py +++ b/coldpress/utils.py @@ -1,4 +1,4 @@ -# Generated by: Claude Sonnet 4.5 +# Assisted by: Claude Sonnet 4.5 """Shared utility functions for Coldpress generators. This module contains common functions used across JobSet, Kubeflow, and KubeRay generators diff --git a/coldpress_setup/__init__.py b/coldpress_setup/__init__.py index 7d51242..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.1" diff --git a/coldpress_setup/cli.py b/coldpress_setup/cli.py index 86708f4..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 diff --git a/coldpress_setup/generator.py b/coldpress_setup/generator.py index 85e454e..b7250ee 100644 --- a/coldpress_setup/generator.py +++ b/coldpress_setup/generator.py @@ -1,4 +1,4 @@ -# Generated by: Claude Sonnet 4.5 +# Assisted by: Claude Sonnet 4.5 """Generate Kubernetes manifests for Coldpress cluster setup.""" import yaml