srunx 0.1.0__tar.gz

srunx-0.1.0/.gitignore

@@ -0,0 +1,12 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+
+.claude/

srunx-0.1.0/.pre-commit-config.yaml

@@ -0,0 +1,43 @@
+repos:
+  # Ruff (linter + formatter)
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.11.13
+    hooks:
+      - id: ruff
+        args: [--fix]
+        require_serial: true
+      - id: ruff-format
+        require_serial: true
+
+  # Pre-commit default hooks
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v5.0.0
+    hooks:
+      - id: check-added-large-files
+      - id: check-json
+      - id: check-toml
+      - id: check-xml
+      - id: check-yaml
+      - id: debug-statements
+      - id: detect-aws-credentials
+        args: [--allow-missing-credentials]
+      - id: detect-private-key
+
+  # uv pip-compile
+  - repo: https://github.com/astral-sh/uv-pre-commit
+    rev: 0.7.13
+    hooks:
+      - id: pip-compile
+        name: pip-compile requirements.in
+        args: [requirements.in, -o, requirements.txt]
+      - id: pip-compile
+        name: pip-compile requirements-dev.in
+        args: [requirements-dev.in, -o, requirements-dev.txt]
+        files: ^requirements-dev\\.(in|txt)$
+
+  # MyPy
+  # - repo: https://github.com/pre-commit/mirrors-mypy
+  #   rev: v1.10.0
+  #   hooks:
+  #     - id: mypy
+  #       args: [--config-file=pyproject.toml]

srunx-0.1.0/.python-version

@@ -0,0 +1 @@
+3.11

srunx-0.1.0/CLAUDE.md

@@ -0,0 +1,139 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Development Commands
+
+### Package Management
+- `uv sync` - Install dependencies
+- `uv add <package>` - Add new dependency
+- `uv run <command>` - Run commands in virtual environment
+
+### CLI Usage
+- `uv run srunx submit <command>` - Submit SLURM job
+- `uv run srunx status <job_id>` - Check job status
+- `uv run srunx list` - List jobs
+- `uv run srunx cancel <job_id>` - Cancel job
+- `uv run srunx flow run <yaml_file>` - Execute workflow from YAML
+- `uv run srunx flow validate <yaml_file>` - Validate workflow YAML
+
+### Testing
+- `uv run pytest` - Run all tests
+- `uv run pytest --cov=srunx` - Run tests with coverage
+- `uv run pytest tests/test_models.py` - Run specific test file
+- `uv run pytest -v` - Run tests with verbose output
+
+### Direct Usage Examples
+- `uv run srunx submit python train.py --name ml_job --gpus-per-node 1`
+- `uv run srunx submit python process.py --conda ml_env --nodes 2`
+- `uv run srunx flow run workflow.yaml`
+
+## Architecture Overview
+
+### New Modular Structure
+```
+src/srunx/
+├── models.py          # Data models and validation
+├── client.py          # SLURM client for job operations
+├── workflows/         # Workflow management
+│   ├── runner.py      # Workflow execution engine
+│   └── tasks.py       # Prefect task definitions
+├── cli/               # Command-line interfaces
+│   ├── main.py        # Main CLI commands
+│   └── workflow.py    # Workflow CLI
+└── templates/         # SLURM script templates
+    ├── base.slurm.jinja
+    └── advanced.slurm.jinja
+```
+
+### Core Components
+
+#### Models (`models.py`)
+- **Job**: Complete job configuration with resources and environment
+- **JobResource**: Resource allocation (nodes, GPUs, memory, time)
+- **JobEnvironment**: Environment setup (conda, venv, sqsh)
+- **JobStatus**: Job status enumeration
+- **Workflow/WorkflowTask**: Workflow definitions with dependencies
+- **render_job_script()**: Template rendering function
+
+#### Client (`client.py`)
+- **Slurm**: Main interface for SLURM operations
+  - `submit_job()`: Submit jobs with full configuration
+  - `retrieve_job()`: Query job status
+  - `cancel_job()`: Cancel running jobs
+  - `list_jobs()`: List user jobs
+  - `wait_for_completion()`: Block until job completes
+
+#### Workflows (`workflows/`)
+- **WorkflowRunner**: YAML workflow execution with Prefect
+- **Prefect Tasks**:
+  - `submit_and_monitor_job()`: Complete job lifecycle
+  - `submit_job_async()`: Async job submission
+  - `wait_for_job()`: Wait for completion
+
+#### CLI (`cli/`)
+- **Main CLI**: Job management commands (submit, status, list, cancel)
+- **Workflow CLI**: YAML workflow execution with validation
+
+### Template System
+- Enhanced Jinja2 templates with conditional resource allocation
+- `base.slurm.jinja`: Simple job template
+- `advanced.slurm.jinja`: Full-featured template with all options
+- Automatic environment setup integration
+
+### Workflow Definition
+Enhanced YAML workflow format:
+```yaml
+name: ml_pipeline
+tasks:
+  - name: preprocess
+    command: ["python", "preprocess.py"]
+    nodes: 1
+
+  - name: train
+    command: ["python", "train.py"]
+    depends_on: [preprocess]
+    gpus_per_node: 1
+    conda: ml_env
+    memory_per_node: "32GB"
+    time_limit: "4:00:00"
+
+  - name: evaluate
+    command: ["python", "evaluate.py"]
+    depends_on: [train]
+    async: true
+```
+
+### Key Improvements
+- **Unified Job Model**: Single `Job` class with comprehensive configuration
+- **Modular Architecture**: Clear separation of concerns
+- **Enhanced CLI**: Subcommands with rich options
+- **Better Error Handling**: Comprehensive validation and error messages
+- **Resource Management**: Full SLURM resource specification
+- **Workflow Validation**: Dependency checking and cycle detection
+
+## Dependencies
+- **Jinja2**: Template rendering
+- **Pydantic**: Data validation and serialization
+- **Prefect**: Workflow orchestration
+- **Loguru**: Structured logging
+- **PyYAML**: YAML parsing
+
+## Code Quality and Linting
+
+### Quality Checks
+- `uv run mypy .` - Type checking with mypy
+- `uv run ruff check .` - Code linting
+- `uv run ruff format .` - Code formatting
+
+### Pre-commit Quality Checks
+Always run these before committing:
+```bash
+uv run pytest && uv run mypy . && uv run ruff check .
+```
+
+# important-instruction-reminders
+Do what has been asked; nothing more, nothing less.
+NEVER create files unless they're absolutely necessary for achieving your goal.
+ALWAYS prefer editing an existing file to creating a new one.
+NEVER proactively create documentation files (*.md) or README files. Only create documentation files if explicitly requested by the User.

srunx-0.1.0/PKG-INFO

@@ -0,0 +1,405 @@
+Metadata-Version: 2.4
+Name: srunx
+Version: 0.1.0
+Summary: Add your description here
+Author-email: ksterx <kostonerx@gmail.com>
+Requires-Python: >=3.11
+Requires-Dist: jinja2>=3.1.6
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: prefect>=3.4.6
+Requires-Dist: pydantic>=2.11.5
+Requires-Dist: pyyaml>=6.0.2
+Description-Content-Type: text/markdown
+
+# srunx
+
+[![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg)](https://www.python.org/downloads/)
+[![Type Checked](https://img.shields.io/badge/type--checked-mypy-blue.svg)](https://mypy.readthedocs.io/)
+[![Code Style](https://img.shields.io/badge/code%20style-ruff-black.svg)](https://github.com/astral-sh/ruff)
+
+A modern Python library for SLURM workload manager integration with workflow orchestration capabilities.
+
+## Features
+
+- 🚀 **Simple Job Submission**: Easy-to-use API for submitting SLURM jobs
+- ⚙️ **Flexible Configuration**: Support for various environments (conda, venv, sqsh)
+- 📋 **Job Management**: Submit, monitor, cancel, and list jobs
+- 🧩 **Workflow Orchestration**: YAML-based workflow definitions with Prefect integration
+- 📝 **Template System**: Customizable Jinja2 templates for SLURM scripts
+- 🛡️ **Type Safe**: Full type hints and mypy compatibility
+- 🖥️ **CLI Tools**: Command-line interfaces for both job management and workflows
+
+## Installation
+
+### Using uv (Recommended)
+
+```bash
+uv add srunx
+```
+
+### Using pip
+
+```bash
+pip install srunx
+```
+
+### Development Installation
+
+```bash
+git clone https://github.com/your-username/srunx.git
+cd srunx
+uv sync --dev
+```
+
+## Quick Start
+
+### Basic Job Submission
+
+```python
+from srunx import Job, JobResource, JobEnvironment, Slurm
+
+# Create a job configuration
+job = Job(
+    name="my_training_job",
+    command=["python", "train.py", "--epochs", "100"],
+    resources=JobResource(
+        nodes=1,
+        gpus_per_node=1,
+        memory_per_node="32GB",
+        time_limit="4:00:00"
+    ),
+    environment=JobEnvironment(conda="ml_env")
+)
+
+# Submit the job
+client = Slurm()
+job = client.run(job)
+print(f"Submitted job {job.job_id}")
+
+# Monitor job status
+job = client.retrieve(job.job_id)
+print(f"Job status: {job.status}")
+```
+
+### Command Line Usage
+
+#### Submit a Job
+
+```bash
+# Basic job submission
+srunx submit python train.py --name ml_job
+
+# With resource specifications
+srunx submit python train.py \
+  --name gpu_job \
+  --gpus-per-node 2 \
+  --memory 64GB \
+  --time 8:00:00
+
+# With environment setup
+srunx submit python train.py \
+  --conda ml_env \
+  --module cuda/11.8 \
+  --module gcc/9.3.0
+```
+
+#### Job Management
+
+```bash
+# Check job status
+srunx status 12345
+
+# List all jobs
+srunx list
+
+# Cancel a job
+srunx cancel 12345
+```
+
+### Workflow Orchestration
+
+Create a workflow YAML file:
+
+```yaml
+# workflow.yaml
+name: ml_pipeline
+tasks:
+  - name: preprocess
+    command: ["python", "preprocess.py"]
+    nodes: 1
+    memory_per_node: "16GB"
+
+  - name: train
+    command: ["python", "train.py"]
+    depends_on: [preprocess]
+    nodes: 1
+    gpus_per_node: 2
+    memory_per_node: "32GB"
+    time_limit: "8:00:00"
+    conda: ml_env
+
+  - name: evaluate
+    command: ["python", "evaluate.py"]
+    depends_on: [train]
+    nodes: 1
+
+  - name: notify
+    command: ["python", "notify.py"]
+    depends_on: [train, evaluate]
+    async: true
+```
+
+Execute the workflow:
+
+```bash
+# Run workflow
+srunx flow workflow.yaml
+
+# Validate workflow without execution
+srunx flow workflow.yaml --validate-only
+
+# Show execution plan
+srunx flow workflow.yaml --dry-run
+```
+
+## Advanced Usage
+
+### Custom Templates
+
+Create a custom SLURM template:
+
+```bash
+#!/bin/bash
+#SBATCH --job-name={{ job_name }}
+#SBATCH --nodes={{ nodes }}
+{% if gpus_per_node > 0 -%}
+#SBATCH --gpus-per-node={{ gpus_per_node }}
+{% endif -%}
+#SBATCH --time={{ time_limit }}
+#SBATCH --output={{ log_dir }}/%x_%j.out
+
+{{ environment_setup }}
+
+srun {{ command }}
+```
+
+Use it with your job:
+
+```python
+job = client.run(job, template_path="custom_template.slurm.jinja")
+```
+
+### Environment Configuration
+
+#### Conda Environment
+
+```python
+environment = JobEnvironment(
+    conda="my_env",
+    env_vars={"CUDA_VISIBLE_DEVICES": "0,1"}
+)
+```
+
+#### SquashFS Images
+
+```python
+environment = JobEnvironment(
+    sqsh="/path/to/container.sqsh",
+    env_vars={"OMP_NUM_THREADS": "8"}
+)
+```
+
+### Programmatic Workflow Execution
+
+```python
+from srunx.workflows import WorkflowRunner
+
+runner = WorkflowRunner()
+workflow = runner.load_from_yaml("workflow.yaml")
+results = runner.execute_workflow(workflow)
+
+print("Job IDs:")
+for task_name, job_id in results.items():
+    print(f"  {task_name}: {job_id}")
+```
+
+### Async Job Submission
+
+```python
+# Submit job without waiting
+job = client.run(job)
+
+# Later, wait for completion
+completed_job = client.monitor(job, poll_interval=30)
+print(f"Job completed with status: {completed_job.status}")
+```
+
+## API Reference
+
+### Core Classes
+
+#### `Job`
+Main job configuration class with resources and environment settings.
+
+#### `JobResource`
+Resource allocation specification (nodes, GPUs, memory, time).
+
+#### `JobEnvironment`
+Environment setup (conda, venv, sqsh, environment variables).
+
+#### `Slurm`
+Main interface for SLURM operations (submit, status, cancel, list).
+
+#### `WorkflowRunner`
+Workflow execution engine with YAML support.
+
+### CLI Commands
+
+#### Main CLI (`srunx`)
+- `submit` - Submit SLURM jobs
+- `status` - Check job status
+- `list` - List jobs
+- `cancel` - Cancel jobs
+
+#### Workflow CLI (`srunx workflow`)
+- Execute YAML-defined workflows
+- Validate workflow files
+- Show execution plans
+
+## Configuration
+
+### Environment Variables
+
+- `SLURM_LOG_DIR`: Default directory for SLURM logs (default: `logs`)
+
+### Template Locations
+
+srunx includes built-in templates:
+- `base.slurm.jinja`: Basic job template
+- `advanced.slurm.jinja`: Full-featured template with all options
+
+## Development
+
+### Setup Development Environment
+
+```bash
+git clone https://github.com/your-username/srunx.git
+cd srunx
+uv sync --dev
+```
+
+### Run Tests
+
+```bash
+uv run pytest
+```
+
+### Type Checking
+
+```bash
+uv run mypy .
+```
+
+### Code Formatting
+
+```bash
+uv run ruff check .
+uv run ruff format .
+```
+
+## Examples
+
+### Machine Learning Pipeline
+
+```python
+# Complete ML pipeline example
+from srunx import Job, JobResource, JobEnvironment, Slurm
+
+def create_ml_job(script: str, **kwargs) -> Job:
+    return Job(
+        name=f"ml_{script.replace('.py', '')}",
+        command=["python", script] + [f"--{k}={v}" for k, v in kwargs.items()],
+        resources=JobResource(
+            nodes=1,
+            gpus_per_node=1,
+            memory_per_node="32GB",
+            time_limit="4:00:00"
+        ),
+        environment=JobEnvironment(conda="pytorch")
+    )
+
+client = Slurm()
+
+# Submit preprocessing job
+prep_job = create_ml_job("preprocess.py", data_path="/data", output_path="/processed")
+prep_job = client.run(prep_job)
+
+# Wait for preprocessing to complete
+client.monitor(prep_job)
+
+# Submit training job
+train_job = create_ml_job("train.py", data_path="/processed", model_path="/models")
+train_job = client.run(train_job)
+
+print(f"Training job {train_job.job_id} submitted")
+```
+
+### Distributed Computing
+
+```python
+# Multi-node distributed job
+distributed_job = Job(
+    name="distributed_training",
+    command=[
+        "mpirun", "-np", "16",
+        "python", "distributed_train.py"
+    ],
+    resources=JobResource(
+        nodes=4,
+        ntasks_per_node=4,
+        cpus_per_task=8,
+        gpus_per_node=2,
+        memory_per_node="128GB",
+        time_limit="12:00:00"
+    ),
+    environment=JobEnvironment(
+        conda="distributed_ml"
+    )
+)
+
+job = client.run(distributed_job)
+```
+
+## Contributing
+
+We welcome contributions! Please see our [Contributing Guide](CONTRIBUTING.md) for details.
+
+### Development Workflow
+
+1. Fork the repository
+2. Create a feature branch
+3. Make your changes
+4. Add tests
+5. Run type checking and tests
+6. Submit a pull request
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
+
+## Changelog
+
+See [CHANGELOG.md](CHANGELOG.md) for release history.
+
+## Support
+
+- 📖 **Documentation**: [docs.example.com/srunx](https://docs.example.com/srunx)
+- 🐞 **Issues**: [GitHub Issues](https://github.com/your-username/srunx/issues)
+- 💬 **Discussions**: [GitHub Discussions](https://github.com/your-username/srunx/discussions)
+
+## Acknowledgments
+
+- Built with [Pydantic](https://pydantic.dev/) for data validation
+- Workflow orchestration powered by [Prefect](https://www.prefect.io/)
+- Template rendering with [Jinja2](https://jinja.palletsprojects.com/)
+- Package management with [uv](https://github.com/astral-sh/uv)