flexium 0.1.0__tar.gz

flexium-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Flexium Authors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

flexium-0.1.0/PKG-INFO

@@ -0,0 +1,328 @@
+Metadata-Version: 2.4
+Name: flexium
+Version: 0.1.0
+Summary: Flexium.AI - Flexible Resource Allocation for GPU training with live migration
+Author: Flexium.AI
+License: MIT
+Project-URL: Homepage, https://github.com/flexiumai/flexium
+Project-URL: Documentation, https://github.com/flexiumai/flexium#readme
+Project-URL: Repository, https://github.com/flexiumai/flexium
+Keywords: gpu,pytorch,training,orchestration,migration,cuda,flexible-resource-allocation
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: System :: Distributed Computing
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: grpcio>=1.50.0
+Requires-Dist: protobuf>=4.0.0
+Requires-Dist: pynvml>=11.0.0
+Requires-Dist: flask>=2.0.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
+Requires-Dist: mypy>=1.0.0; extra == "dev"
+Requires-Dist: ruff>=0.1.0; extra == "dev"
+Requires-Dist: grpcio-tools>=1.50.0; extra == "dev"
+Provides-Extra: torch
+Requires-Dist: torch>=1.9.0; extra == "torch"
+Requires-Dist: torchvision>=0.10.0; extra == "torch"
+Provides-Extra: torch-cuda
+Requires-Dist: pytorch-lightning>=2.0.0; extra == "torch-cuda"
+Provides-Extra: lightning
+Requires-Dist: pytorch-lightning>=2.0.0; extra == "lightning"
+Provides-Extra: debug
+Requires-Dist: debugpy>=1.6.0; extra == "debug"
+Provides-Extra: all
+Requires-Dist: flexium[debug,dev,torch]; extra == "all"
+Dynamic: license-file
+
+<p align="center">
+  <img src="logo_with_text.png" alt="Flexium.AI Logo" width="400">
+</p>
+
+<h1 align="center">Flexium.AI</h1>
+
+<p align="center">
+  <strong>Flexible Resource Allocation for GPU Training</strong><br>
+  Live GPU migration for PyTorch - zero code changes required.
+</p>
+
+<p align="center">
+  <a href="#installation">Installation</a> •
+  <a href="#quick-start">Quick Start</a> •
+  <a href="#features">Features</a> •
+  <a href="#documentation">Documentation</a>
+</p>
+
+---
+
+Flexium.AI allows you to move running PyTorch training jobs between GPUs without stopping them. Free up GPUs for colleagues, recover from hardware errors, and manage shared GPU clusters - all while your training continues seamlessly.
+
+## The Problem
+
+When you need to free a GPU during training, your options are bad:
+
+```python
+# You're training on cuda:0, colleague needs it urgently
+model = model.to("cuda:1")  # Doesn't work - memory stays on cuda:0!
+torch.cuda.empty_cache()     # Still doesn't free it
+# Your colleague still can't use the GPU
+```
+
+PyTorch's CUDA memory allocator holds onto GPU memory even after `model.to()`. The only way to truly free a GPU is to kill the process - losing your training progress.
+
+## The Solution
+
+Flexium guarantees complete GPU release through driver-level migration (requires NVIDIA driver 580+):
+
+```python
+import flexium.auto  # Add this import
+
+with flexium.auto.run():  # Wrap your training
+    model = Net().cuda()  # Standard PyTorch - no changes needed!
+    for epoch in range(100):
+        for batch in dataloader:
+            # ... your normal training code ...
+```
+
+Now you can migrate anytime:
+
+```bash
+flexium-ctl migrate <process-id> cuda:1  # Training continues on new GPU
+```
+
+The GPU is **100% freed** - your colleague can use it immediately.
+
+## Features
+
+- **Zero code changes** - Just wrap your training in `flexium.auto.run()`
+- **True GPU freedom** - Guarantees zero memory residue on source GPU
+- **Live migration** - Move training between GPUs without losing progress
+- **Automatic state management** - Training state preserved transparently
+- **Web dashboard** - Visual GPU management at `http://localhost:8080`
+- **CLI control** - `flexium-ctl` for scripting and automation
+- **Pause/Resume** - Free GPU completely, resume later on any available GPU
+- **Error recovery** - Auto-migrate on OOM, ECC errors, or hardware failures
+- **Graceful degradation** - Training continues even if orchestrator dies
+
+## Installation
+
+```bash
+# From PyPI (coming soon)
+pip install flexium
+
+# From source
+git clone https://github.com/your-org/flexium.git
+cd flexium
+pip install -e .
+```
+
+### Requirements
+
+- Python 3.8+
+- PyTorch 2.0+ with CUDA support ([install guide](https://pytorch.org/get-started/locally/))
+- NVIDIA GPU with CUDA support
+- **NVIDIA Driver 580+** (required for zero-residue migration)
+- Linux x86_64 (Windows/macOS not yet supported)
+
+## Quick Start
+
+### 1. Start the Orchestrator
+
+```bash
+flexium-ctl server --dashboard
+```
+
+This starts:
+- gRPC server on port 50051 (coordinates migrations)
+- Web dashboard at http://localhost:8080
+
+### 2. Add Flexium to Your Training
+
+```python
+import flexium.auto
+import torch
+
+with flexium.auto.run():
+    model = MyModel().cuda()
+    optimizer = torch.optim.Adam(model.parameters())
+
+    for epoch in range(100):
+        for batch in dataloader:
+            loss = model(batch.cuda()).mean()
+            loss.backward()
+            optimizer.step()
+```
+
+That's it! Your training is now migratable.
+
+### 3. Migrate When Needed
+
+**Via Dashboard:**
+Open http://localhost:8080, find your process, click "Migrate", select target GPU.
+
+**Via CLI:**
+```bash
+# List running processes
+flexium-ctl list
+
+# Migrate to different GPU
+flexium-ctl migrate <process-id> cuda:1
+
+# Pause (free GPU completely)
+flexium-ctl pause <process-id>
+
+# Resume on any available GPU
+flexium-ctl resume <process-id>
+```
+
+## Real-World Scenarios
+
+### GPU Contention
+```bash
+# Alice is training on cuda:0, Bob needs it urgently
+flexium-ctl migrate alice-abc123 cuda:2
+# Alice's training continues on cuda:2, cuda:0 is free for Bob
+```
+
+### OOM Recovery
+```python
+with flexium.auto.run():
+    # If OOM occurs, flexium auto-migrates to GPU with more VRAM
+    model = LargeModel().cuda()
+    train(model)
+```
+
+### Shared Cluster Management
+```bash
+# See all jobs across the cluster
+flexium-ctl list
+
+# Dashboard shows: who's using what, progress, VRAM usage
+# One-click migration when someone needs a GPU
+```
+
+### Hardware Failure
+```python
+with flexium.auto.run():
+    # ECC error? Driver crash? Auto-migrate to healthy GPU
+    # Bad GPU marked unhealthy, won't be used again
+    train(model)
+```
+
+## How It Works
+
+Flexium uses **driver-level GPU migration** (requires NVIDIA driver 580+) that guarantees complete GPU memory release - something that's impossible with standard PyTorch memory management.
+
+When you request a migration:
+1. GPU state is checkpointed at the driver level
+2. Source GPU memory is completely released (0 MB residue)
+3. State is restored on the target GPU
+4. Training continues seamlessly
+
+The orchestrator coordinates all GPU resources across your cluster, tracking utilization, health status, and process state to enable intelligent scheduling and automatic recovery.
+
+## Documentation
+
+- [Getting Started Guide](docs/getting-started.md)
+- [API Reference](docs/api.md)
+- [Architecture Overview](docs/ARCHITECTURE.md)
+- [Examples](docs/examples.md)
+- [Troubleshooting](docs/troubleshooting.md)
+
+## Benchmarks
+
+Migration overhead is minimal:
+
+| Model Size | Checkpoint Time | Total Migration |
+|------------|-----------------|-----------------|
+| 100 MB     | ~200 ms         | ~700 ms         |
+| 500 MB     | ~800 ms         | ~2 sec          |
+| 1 GB       | ~1.5 sec        | ~4 sec          |
+
+Runtime overhead during normal training: **< 2%**
+
+## Configuration
+
+### Environment Variables
+```bash
+export GPU_ORCHESTRATOR=localhost:50051
+export GPU_DEVICE=cuda:0
+```
+
+### Config File (`~/.flexiumrc`)
+```yaml
+orchestrator: localhost:50051
+device: cuda:0
+```
+
+### Inline Parameters
+```python
+with flexium.auto.run(orchestrator="server:50051", device="cuda:1"):
+    ...
+```
+
+## CLI Reference
+
+```bash
+# Server management
+flexium-ctl server [--port PORT] [--dashboard]
+
+# Process management
+flexium-ctl list                          # List all processes
+flexium-ctl status <process-id>           # Get process details
+flexium-ctl migrate <process-id> <device> # Migrate to device
+flexium-ctl pause <process-id>            # Pause (free GPU)
+flexium-ctl resume <process-id> [device]  # Resume training
+
+# GPU management
+flexium-ctl devices                       # List all GPUs
+flexium-ctl device <gpu-uuid>             # Get GPU details
+```
+
+## Contributing
+
+We welcome contributions! Please see our contributing guidelines (coming soon).
+
+```bash
+# Development setup
+git clone https://github.com/your-org/flexium.git
+cd flexium
+pip install -e ".[dev]"
+
+# Run tests
+pytest tests/ -v
+
+# Run linting
+ruff check flexium/
+mypy flexium/
+```
+
+## License
+
+MIT License - see [LICENSE](LICENSE) for details.
+
+## Acknowledgments
+
+Flexium.AI was inspired by the challenges of managing shared GPU clusters in research environments. Special thanks to everyone who provided feedback and tested early versions.
+
+---
+
+<p align="center">
+  <img src="logo_with_text.png" alt="Flexium.AI" width="150"><br>
+  <strong>Flexium.AI</strong> - Flexible Resource Allocation<br>
+  <em>Making GPU management effortless</em>
+</p>
+
+**Note:** Flexium.AI is currently in alpha. APIs may change. Please report issues on GitHub.

flexium-0.1.0/README.md

@@ -0,0 +1,279 @@
+<p align="center">
+  <img src="logo_with_text.png" alt="Flexium.AI Logo" width="400">
+</p>
+
+<h1 align="center">Flexium.AI</h1>
+
+<p align="center">
+  <strong>Flexible Resource Allocation for GPU Training</strong><br>
+  Live GPU migration for PyTorch - zero code changes required.
+</p>
+
+<p align="center">
+  <a href="#installation">Installation</a> •
+  <a href="#quick-start">Quick Start</a> •
+  <a href="#features">Features</a> •
+  <a href="#documentation">Documentation</a>
+</p>
+
+---
+
+Flexium.AI allows you to move running PyTorch training jobs between GPUs without stopping them. Free up GPUs for colleagues, recover from hardware errors, and manage shared GPU clusters - all while your training continues seamlessly.
+
+## The Problem
+
+When you need to free a GPU during training, your options are bad:
+
+```python
+# You're training on cuda:0, colleague needs it urgently
+model = model.to("cuda:1")  # Doesn't work - memory stays on cuda:0!
+torch.cuda.empty_cache()     # Still doesn't free it
+# Your colleague still can't use the GPU
+```
+
+PyTorch's CUDA memory allocator holds onto GPU memory even after `model.to()`. The only way to truly free a GPU is to kill the process - losing your training progress.
+
+## The Solution
+
+Flexium guarantees complete GPU release through driver-level migration (requires NVIDIA driver 580+):
+
+```python
+import flexium.auto  # Add this import
+
+with flexium.auto.run():  # Wrap your training
+    model = Net().cuda()  # Standard PyTorch - no changes needed!
+    for epoch in range(100):
+        for batch in dataloader:
+            # ... your normal training code ...
+```
+
+Now you can migrate anytime:
+
+```bash
+flexium-ctl migrate <process-id> cuda:1  # Training continues on new GPU
+```
+
+The GPU is **100% freed** - your colleague can use it immediately.
+
+## Features
+
+- **Zero code changes** - Just wrap your training in `flexium.auto.run()`
+- **True GPU freedom** - Guarantees zero memory residue on source GPU
+- **Live migration** - Move training between GPUs without losing progress
+- **Automatic state management** - Training state preserved transparently
+- **Web dashboard** - Visual GPU management at `http://localhost:8080`
+- **CLI control** - `flexium-ctl` for scripting and automation
+- **Pause/Resume** - Free GPU completely, resume later on any available GPU
+- **Error recovery** - Auto-migrate on OOM, ECC errors, or hardware failures
+- **Graceful degradation** - Training continues even if orchestrator dies
+
+## Installation
+
+```bash
+# From PyPI (coming soon)
+pip install flexium
+
+# From source
+git clone https://github.com/your-org/flexium.git
+cd flexium
+pip install -e .
+```
+
+### Requirements
+
+- Python 3.8+
+- PyTorch 2.0+ with CUDA support ([install guide](https://pytorch.org/get-started/locally/))
+- NVIDIA GPU with CUDA support
+- **NVIDIA Driver 580+** (required for zero-residue migration)
+- Linux x86_64 (Windows/macOS not yet supported)
+
+## Quick Start
+
+### 1. Start the Orchestrator
+
+```bash
+flexium-ctl server --dashboard
+```
+
+This starts:
+- gRPC server on port 50051 (coordinates migrations)
+- Web dashboard at http://localhost:8080
+
+### 2. Add Flexium to Your Training
+
+```python
+import flexium.auto
+import torch
+
+with flexium.auto.run():
+    model = MyModel().cuda()
+    optimizer = torch.optim.Adam(model.parameters())
+
+    for epoch in range(100):
+        for batch in dataloader:
+            loss = model(batch.cuda()).mean()
+            loss.backward()
+            optimizer.step()
+```
+
+That's it! Your training is now migratable.
+
+### 3. Migrate When Needed
+
+**Via Dashboard:**
+Open http://localhost:8080, find your process, click "Migrate", select target GPU.
+
+**Via CLI:**
+```bash
+# List running processes
+flexium-ctl list
+
+# Migrate to different GPU
+flexium-ctl migrate <process-id> cuda:1
+
+# Pause (free GPU completely)
+flexium-ctl pause <process-id>
+
+# Resume on any available GPU
+flexium-ctl resume <process-id>
+```
+
+## Real-World Scenarios
+
+### GPU Contention
+```bash
+# Alice is training on cuda:0, Bob needs it urgently
+flexium-ctl migrate alice-abc123 cuda:2
+# Alice's training continues on cuda:2, cuda:0 is free for Bob
+```
+
+### OOM Recovery
+```python
+with flexium.auto.run():
+    # If OOM occurs, flexium auto-migrates to GPU with more VRAM
+    model = LargeModel().cuda()
+    train(model)
+```
+
+### Shared Cluster Management
+```bash
+# See all jobs across the cluster
+flexium-ctl list
+
+# Dashboard shows: who's using what, progress, VRAM usage
+# One-click migration when someone needs a GPU
+```
+
+### Hardware Failure
+```python
+with flexium.auto.run():
+    # ECC error? Driver crash? Auto-migrate to healthy GPU
+    # Bad GPU marked unhealthy, won't be used again
+    train(model)
+```
+
+## How It Works
+
+Flexium uses **driver-level GPU migration** (requires NVIDIA driver 580+) that guarantees complete GPU memory release - something that's impossible with standard PyTorch memory management.
+
+When you request a migration:
+1. GPU state is checkpointed at the driver level
+2. Source GPU memory is completely released (0 MB residue)
+3. State is restored on the target GPU
+4. Training continues seamlessly
+
+The orchestrator coordinates all GPU resources across your cluster, tracking utilization, health status, and process state to enable intelligent scheduling and automatic recovery.
+
+## Documentation
+
+- [Getting Started Guide](docs/getting-started.md)
+- [API Reference](docs/api.md)
+- [Architecture Overview](docs/ARCHITECTURE.md)
+- [Examples](docs/examples.md)
+- [Troubleshooting](docs/troubleshooting.md)
+
+## Benchmarks
+
+Migration overhead is minimal:
+
+| Model Size | Checkpoint Time | Total Migration |
+|------------|-----------------|-----------------|
+| 100 MB     | ~200 ms         | ~700 ms         |
+| 500 MB     | ~800 ms         | ~2 sec          |
+| 1 GB       | ~1.5 sec        | ~4 sec          |
+
+Runtime overhead during normal training: **< 2%**
+
+## Configuration
+
+### Environment Variables
+```bash
+export GPU_ORCHESTRATOR=localhost:50051
+export GPU_DEVICE=cuda:0
+```
+
+### Config File (`~/.flexiumrc`)
+```yaml
+orchestrator: localhost:50051
+device: cuda:0
+```
+
+### Inline Parameters
+```python
+with flexium.auto.run(orchestrator="server:50051", device="cuda:1"):
+    ...
+```
+
+## CLI Reference
+
+```bash
+# Server management
+flexium-ctl server [--port PORT] [--dashboard]
+
+# Process management
+flexium-ctl list                          # List all processes
+flexium-ctl status <process-id>           # Get process details
+flexium-ctl migrate <process-id> <device> # Migrate to device
+flexium-ctl pause <process-id>            # Pause (free GPU)
+flexium-ctl resume <process-id> [device]  # Resume training
+
+# GPU management
+flexium-ctl devices                       # List all GPUs
+flexium-ctl device <gpu-uuid>             # Get GPU details
+```
+
+## Contributing
+
+We welcome contributions! Please see our contributing guidelines (coming soon).
+
+```bash
+# Development setup
+git clone https://github.com/your-org/flexium.git
+cd flexium
+pip install -e ".[dev]"
+
+# Run tests
+pytest tests/ -v
+
+# Run linting
+ruff check flexium/
+mypy flexium/
+```
+
+## License
+
+MIT License - see [LICENSE](LICENSE) for details.
+
+## Acknowledgments
+
+Flexium.AI was inspired by the challenges of managing shared GPU clusters in research environments. Special thanks to everyone who provided feedback and tested early versions.
+
+---
+
+<p align="center">
+  <img src="logo_with_text.png" alt="Flexium.AI" width="150"><br>
+  <strong>Flexium.AI</strong> - Flexible Resource Allocation<br>
+  <em>Making GPU management effortless</em>
+</p>
+
+**Note:** Flexium.AI is currently in alpha. APIs may change. Please report issues on GitHub.

flexium-0.1.0/flexium/__init__.py

@@ -0,0 +1,39 @@
+"""Flexium - Dynamic GPU orchestration for PyTorch training.
+
+This package provides GPU orchestration with live migration support using
+driver-level migration for zero-residue GPU migration (requires driver 580+).
+
+Usage:
+
+    import flexium.auto
+
+    with flexium.auto.run():
+        model = nn.Linear(784, 10).cuda()
+        optimizer = torch.optim.Adam(model.parameters())
+
+        for epoch in range(100):
+            for batch in dataloader:
+                data, target = batch[0].cuda(), batch[1].cuda()
+                loss = model(data).sum()
+                loss.backward()
+                optimizer.step()
+
+Configuration:
+    - Set GPU_ORCHESTRATOR environment variable to orchestrator address
+    - Or pass orchestrator= parameter to flexium.auto.run()
+
+Migration is transparent - training continues in the same process,
+same loop iteration, just on a different GPU.
+"""
+
+from __future__ import annotations
+
+__version__ = "0.1.0"
+
+# The main API is flexium.auto.run()
+# Import with: import flexium.auto
+# Usage: with flexium.auto.run(): ...
+
+__all__ = [
+    "__version__",
+]