optirs-tpu 0.2.0

# OptiRS TPU

TPU coordination and pod management for large-scale distributed optimization in the OptiRS machine learning optimization library.

## Overview

OptiRS-TPU provides comprehensive support for Google Cloud TPU (Tensor Processing Unit) coordination, pod management, and distributed training optimization. This crate enables efficient scaling of machine learning optimization workloads across TPU pods, with intelligent resource allocation, fault tolerance, and performance monitoring.

## Features

- **TPU Pod Management**: Automatic TPU pod allocation and coordination
- **Distributed Optimization**: Seamless scaling across TPU cores and pods
- **XLA Integration**: Optimized compilation for TPU execution
- **Fault Tolerance**: Robust error handling and recovery mechanisms
- **Dynamic Scaling**: Automatic resource scaling based on workload demands
- **Performance Monitoring**: Real-time TPU utilization and optimization metrics
- **Cost Optimization**: Intelligent resource allocation to minimize costs
- **Multi-Region Support**: Cross-region TPU coordination capabilities

## TPU Architecture Support

### TPU Generations
- **TPU v4**: Latest generation with enhanced performance and memory
- **TPU v3**: High-performance training and inference
- **TPU v2**: Cost-effective training for medium-scale models
- **TPU Edge**: Edge deployment optimization (future support)

### Pod Configurations
- **Single TPU**: Development and small-scale training
- **TPU Pod Slice**: Multi-core coordination (8, 32, 128 cores)
- **Full TPU Pod**: Large-scale distributed training (256+ cores)
- **Multi-Pod**: Cross-pod coordination for massive workloads

## Installation

Add this to your `Cargo.toml`:

```toml
[dependencies]
optirs-tpu = "0.2.0"
scirs2-core = "0.1.1"  # Required foundation
```

### Prerequisites

1. **Google Cloud Setup**:
   ```bash
   gcloud auth login
   gcloud config set project YOUR_PROJECT_ID
   ```

2. **TPU Quotas**: Ensure sufficient TPU quotas in your GCP project

3. **Network Configuration**: Proper VPC and firewall settings for TPU communication

## Usage

### Basic TPU Optimization

```rust
use optirs_tpu::{TpuManager, TpuOptimizer, PodConfig};
use optirs_core::optimizers::Adam;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize TPU manager
    let tpu_manager = TpuManager::new()
        .with_project("your-gcp-project")
        .with_zone("us-central1-a")
        .build()
        .await?;

    // Request TPU resources
    let pod_config = PodConfig::new()
        .with_tpu_type("v3-8")
        .with_cores(8)
        .build();

    let tpu_pod = tpu_manager.allocate_pod(&pod_config).await?;

    // Create TPU-optimized optimizer
    let mut optimizer = TpuOptimizer::new(tpu_pod)
        .with_optimizer(Adam::new(0.001))
        .with_xla_optimization(true)
        .build()?;

    // Distributed optimization across TPU cores
    let mut params = optimizer.create_distributed_tensor(&[8192, 4096])?;
    let grads = optimizer.load_gradients_from_dataset(&training_data).await?;

    // Perform distributed optimization step
    optimizer.step_distributed(&mut params, &grads).await?;

    Ok(())
}
```

### Multi-Pod Training

```rust
use optirs_tpu::{MultiPodManager, PodTopology, AllReduceStrategy};

// Setup multi-pod coordination
let multi_pod = MultiPodManager::new()
    .with_topology(PodTopology::Ring)
    .with_allreduce_strategy(AllReduceStrategy::HierarchicalAllReduce)
    .with_pods(&[
        pod_config_a.in_zone("us-central1-a"),
        pod_config_b.in_zone("us-central1-b"),
        pod_config_c.in_zone("europe-west4-a"),
    ])
    .build()
    .await?;

// Distribute model across pods
multi_pod.distribute_model(&model_parameters).await?;

// Synchronized optimization across all pods
let optimization_result = multi_pod
    .step_synchronized(&gradients)
    .with_timeout(Duration::from_secs(300))
    .await?;
```

### Dynamic Resource Scaling

```rust
use optirs_tpu::{AutoScaler, ScalingPolicy, WorkloadMetrics};

// Setup automatic scaling
let autoscaler = AutoScaler::new()
    .with_policy(ScalingPolicy::PerformanceBased)
    .with_min_cores(8)
    .with_max_cores(256)
    .with_target_utilization(0.85)
    .build();

// Monitor and scale based on workload
let metrics = WorkloadMetrics::from_current_training(&optimizer).await?;
let scaling_decision = autoscaler.evaluate_scaling(&metrics).await?;

if let Some(new_config) = scaling_decision {
    optimizer.scale_to_config(new_config).await?;
}
```

### XLA Optimization

```rust
use optirs_tpu::xla::{XlaCompiler, OptimizationFlags};

// Compile optimization pipeline with XLA
let xla_compiler = XlaCompiler::new()
    .with_optimization_level(3)
    .with_flags(OptimizationFlags {
        enable_fusion: true,
        enable_layout_optimization: true,
        enable_memory_optimization: true,
    })
    .build();

let optimized_graph = xla_compiler
    .compile_optimization_pipeline(&optimizer_config)
    .await?;

optimizer.load_xla_graph(optimized_graph)?;
```

## Architecture

### TPU Resource Management
- **Pod Allocation**: Intelligent TPU pod allocation and deallocation
- **Resource Scheduling**: Optimal scheduling across available TPU resources
- **Cost Management**: Cost-aware resource allocation strategies
- **Health Monitoring**: Continuous TPU health and performance monitoring

### Distributed Coordination
- **Communication Patterns**: Efficient inter-TPU communication protocols
- **Synchronization**: Gradient synchronization across TPU cores and pods
- **Load Balancing**: Dynamic load balancing for optimal utilization
- **Fault Recovery**: Automatic recovery from TPU failures

### Performance Optimization
- **Memory Management**: Efficient TPU memory utilization
- **Computation Overlap**: Overlapping computation and communication
- **Pipeline Parallelism**: Advanced pipelining for large models
- **Mixed Precision**: Automatic mixed precision optimization

## Configuration

### Environment Variables
```bash
export GOOGLE_APPLICATION_CREDENTIALS="path/to/service-account.json"
export TPU_PROJECT="your-gcp-project"
export TPU_ZONE="us-central1-a"
export TPU_NETWORK="default"
```

### Configuration File
```yaml
# tpu_config.yaml
project: "your-gcp-project"
zone: "us-central1-a"
network: "default"
accelerator_type: "v3-8"
software_version: "2.8.0"
enable_ip_alias: true
reserved: false
```

Load configuration:
```rust
use optirs_tpu::config::TpuConfig;

let config = TpuConfig::from_file("tpu_config.yaml").await?;
let tpu_manager = TpuManager::from_config(config).await?;
```

## Monitoring and Observability

### Performance Metrics
- TPU utilization rates
- Memory usage patterns
- Communication overhead
- Computation efficiency
- Power consumption

### Logging Integration
```rust
use optirs_tpu::monitoring::{TpuMonitor, MetricsCollector};

let monitor = TpuMonitor::new()
    .with_collection_interval(Duration::from_secs(30))
    .with_metrics_endpoint("https://monitoring.googleapis.com")
    .build();

// Collect and report TPU metrics
let metrics = monitor.collect_metrics().await?;
monitor.report_to_cloud_monitoring(&metrics).await?;
```

## Error Handling

OptiRS-TPU provides comprehensive error handling for distributed TPU operations:

```rust
use optirs_tpu::error::{TpuError, TpuResult};

match optimizer.step_distributed(&mut params, &grads).await {
    Ok(metrics) => {
        println!("Optimization successful: {:.4} loss", metrics.loss);
    }
    Err(TpuError::PodUnavailable) => {
        // Handle TPU pod unavailability
        tpu_manager.request_alternate_pod().await?;
    }
    Err(TpuError::CommunicationTimeout) => {
        // Handle communication timeouts
        optimizer.reduce_batch_size().await?;
    }
    Err(TpuError::OutOfMemory) => {
        // Handle TPU memory exhaustion
        optimizer.enable_gradient_checkpointing().await?;
    }
    Err(e) => return Err(e.into()),
}
```

## Cost Optimization

### Preemptible TPUs
```rust
use optirs_tpu::preemptible::{PreemptibleManager, PreemptionHandler};

let preemptible_config = PodConfig::new()
    .with_tpu_type("v3-8")
    .with_preemptible(true)  // 70% cost savings
    .build();

let preemption_handler = PreemptionHandler::new()
    .with_checkpoint_frequency(Duration::from_secs(300))
    .with_auto_restart(true)
    .build();

let tpu_pod = tpu_manager
    .allocate_preemptible_pod(&preemptible_config, preemption_handler)
    .await?;
```

### Spot TPU Bidding
```rust
use optirs_tpu::spot::{SpotManager, BiddingStrategy};

let spot_manager = SpotManager::new()
    .with_bidding_strategy(BiddingStrategy::CostOptimized)
    .with_max_price_per_hour(2.50)
    .build();

let spot_pod = spot_manager.bid_for_pod(&pod_config).await?;
```

## Platform Support

| Feature | Google Cloud TPU | TPU Research Cloud | On-Premise |
|---------|------------------|-------------------|-------------|
| v2 TPUs | ✅ | ✅ | ❌ |
| v3 TPUs | ✅ | ✅ | ❌ |
| v4 TPUs | ✅ | ✅ | ❌ |
| Pod Management | ✅ | Limited | ❌ |
| Preemptible | ✅ | ❌ | ❌ |

## Development Guidelines

### Coding Standards

To maintain consistency and readability across the codebase, please follow these guidelines:

#### Variable Naming
- **Always use `snake_case` for variable names** (e.g., `user_id`, `max_iterations`, `learning_rate`)
- **Avoid camelCase or other naming conventions** (e.g., `userId` ❌, `maxIterations` ❌)
- **Use descriptive names** that clearly indicate the variable's purpose

```rust
// ✅ Correct: snake_case
let experiment_id = "exp_001";
let max_epochs = 100;
let learning_rate = 0.001;

// ❌ Incorrect: camelCase or other formats
let experimentId = "exp_001";
let maxEpochs = 100;
let learningrate = 0.001;
```

#### Function and Method Names
- Use `snake_case` for function and method names
- Use descriptive verbs that indicate the function's action

#### Type Names
- Use `PascalCase` for struct, enum, and trait names
- Use `SCREAMING_SNAKE_CASE` for constants

#### General Guidelines
- Follow Rust's official naming conventions as specified in [RFC 430](https://github.com/rust-lang/rfcs/blob/master/text/0430-finalizing-naming-conventions.md)
- Use `rustfmt` and `clippy` to maintain code formatting and catch common issues
- Write clear, self-documenting code with appropriate comments

### Before Submitting Code
1. Run `cargo fmt` to format your code
2. Run `cargo clippy` to check for lint issues
3. Ensure all tests pass with `cargo test`
4. Verify compilation with `cargo check`

## Contributing

OptiRS follows the Cool Japan organization's development standards. See the main OptiRS repository for contribution guidelines.

## License

This project is licensed under the Apache License, Version 2.0.