Stop Cluster

def stop_idle_clusters(idle_threshold_minutes=30):
    """Stop clusters that have been idle for too long"""
    
    # Get all running clusters
    response = requests.get(
        "https://api.tensorone.ai/v1/clusters",
        headers={"Authorization": f"Bearer {API_KEY}"},
        params={"status": "running"}
    )
    
    clusters = response.json()["data"]["clusters"]
    stopped_clusters = []
    
    for cluster in clusters:
        # Check if cluster has been idle
        last_activity = cluster.get("last_accessed")
        if last_activity:
            from datetime import datetime, timedelta
            last_activity_time = datetime.fromisoformat(last_activity.replace('Z', '+00:00'))
            idle_time = datetime.now().timestamp() - last_activity_time.timestamp()
            idle_minutes = idle_time / 60
            
            if idle_minutes > idle_threshold_minutes:
                # Stop idle cluster
                result = stop_cluster_graceful(
                    cluster["id"], 
                    grace_minutes=5,
                    reason=f"Auto-stop: idle for {idle_minutes:.1f} minutes"
                )
                
                if result["success"]:
                    stopped_clusters.append({
                        "cluster_id": cluster["id"],
                        "name": cluster["name"],
                        "idle_minutes": idle_minutes,
                        "cost_saved_per_hour": cluster["cost"]["hourly_rate"]
                    })
    
    return stopped_clusters

async function scheduleMaintenanceStop(clusterIds, maintenanceWindow) {
  const results = [];
  
  for (const clusterId of clusterIds) {
    try {
      // Get cluster info to determine appropriate grace period
      const clusterResponse = await fetch(`https://api.tensorone.ai/v1/clusters/${clusterId}`, {
        headers: { 'Authorization': 'Bearer ' + API_KEY }
      });
      
      const cluster = await clusterResponse.json();
      
      if (!cluster.success) {
        results.push({ clusterId, success: false, error: cluster.error.message });
        continue;
      }
      
      // Determine grace period based on cluster type
      let gracePeriod = 5; // default
      if (cluster.data.name.includes('training')) {
        gracePeriod = 15; // More time for training clusters
      } else if (cluster.data.name.includes('prod')) {
        gracePeriod = 10; // Production clusters need time to drain
      }
      
      const stopResult = await stopClusterGracefully(clusterId, {
        gracePeriod: gracePeriod,
        createSnapshot: true,
        snapshotName: `maintenance_${maintenanceWindow.id}_${clusterId}`,
        reason: `Scheduled maintenance: ${maintenanceWindow.description}`,
        notifyUsers: cluster.data.project_info.team_members || []
      });
      
      results.push({
        clusterId,
        success: true,
        gracePeriod,
        snapshotId: stopResult.snapshot_id,
        estimatedDowntime: maintenanceWindow.duration_minutes
      });
      
    } catch (error) {
      results.push({
        clusterId,
        success: false,
        error: error.message
      });
    }
  }
  
  return results;
}

def handle_training_completion(cluster_id, model_name, final_checkpoint_path):
    """Handle training completion with comprehensive state saving"""
    
    # Create final snapshot with model artifacts
    snapshot_name = f"final_model_{model_name}_{int(time.time())}"
    
    stop_config = {
        "force": False,
        "grace_period_minutes": 20,  # Allow time for final checkpoint
        "save_state": True,
        "snapshot_name": snapshot_name,
        "preserve_data": True,
        "wait_for_completion": True,
        "stop_reason": f"Training completed for model: {model_name}",
        "timeout_minutes": 30
    }
    
    # Add environment variable to signal final checkpoint save
    # This would be picked up by the training script
    update_env_response = requests.patch(
        f"https://api.tensorone.ai/v1/clusters/{cluster_id}/environment",
        headers={"Authorization": f"Bearer {API_KEY}"},
        json={
            "TRAINING_COMPLETED": "true",
            "FINAL_CHECKPOINT_PATH": final_checkpoint_path,
            "MODEL_NAME": model_name
        }
    )
    
    # Wait a bit for the training script to notice and save final checkpoint
    time.sleep(60)
    
    # Now stop the cluster
    response = requests.post(
        f"https://api.tensorone.ai/v1/clusters/{cluster_id}/stop",
        headers={"Authorization": f"Bearer {API_KEY}"},
        json=stop_config
    )
    
    result = response.json()
    
    if result["success"]:
        completion_info = {
            "cluster_id": cluster_id,
            "model_name": model_name,
            "final_snapshot_id": result["data"]["snapshot_creation"]["snapshot_id"],
            "training_duration_hours": result["data"]["cost_summary"]["session_duration_hours"],
            "total_cost": result["data"]["cost_summary"]["total_session_cost"],
            "final_checkpoint_path": final_checkpoint_path,
            "stopped_at": result["data"]["stopped_at"]
        }
        
        # Log completion
        print(f"Training completed for {model_name}")
        print(f"Final snapshot: {completion_info['final_snapshot_id']}")
        print(f"Total cost: ${completion_info['total_cost']:.2f}")
        
        return completion_info
    
    return result

def intelligent_batch_stop(project_id, stop_criteria):
    """Stop clusters based on intelligent criteria"""
    
    # Get all clusters for the project
    response = requests.get(
        "https://api.tensorone.ai/v1/clusters",
        headers={"Authorization": f"Bearer {API_KEY}"},
        params={
            "project_id": project_id,
            "status": "running",
            "include_metrics": True
        }
    )
    
    clusters = response.json()["data"]["clusters"]
    stop_decisions = []
    
    for cluster in clusters:
        decision = {
            "cluster_id": cluster["id"],
            "name": cluster["name"],
            "should_stop": False,
            "reason": "",
            "stop_config": {}
        }
        
        metrics = cluster["metrics"]["current"]
        cost = cluster["cost"]["hourly_rate"]
        
        # Criteria 1: Low utilization + high cost
        if metrics["gpu_utilization"] < stop_criteria.get("min_gpu_util", 20) and cost > stop_criteria.get("max_cost_for_low_util", 5.0):
            decision["should_stop"] = True
            decision["reason"] = f"Low utilization ({metrics['gpu_utilization']}%) with high cost (${cost}/hr)"
            decision["stop_config"] = {
                "grace_period_minutes": 5,
                "save_state": False,
                "reason": "Cost optimization - low utilization"
            }
        
        # Criteria 2: Development clusters after hours
        elif "dev" in cluster["name"].lower() and stop_criteria.get("stop_dev_after_hours", False):
            decision["should_stop"] = True
            decision["reason"] = "Development cluster - after hours shutdown"
            decision["stop_config"] = {
                "grace_period_minutes": 10,
                "save_state": True,
                "reason": "Scheduled dev environment shutdown"
            }
        
        # Criteria 3: Idle clusters
        elif cluster.get("uptime_seconds", 0) > stop_criteria.get("max_idle_seconds", 3600):
            last_access = cluster.get("last_accessed")
            if last_access:
                # Calculate idle time and decide
                pass  # Implementation depends on last access format
        
        stop_decisions.append(decision)
    
    # Execute stop operations
    results = []
    for decision in stop_decisions:
        if decision["should_stop"]:
            result = stop_cluster_graceful(
                decision["cluster_id"],
                **decision["stop_config"]
            )
            results.append({
                "cluster_id": decision["cluster_id"],
                "name": decision["name"],
                "reason": decision["reason"],
                "success": result["success"],
                "response": result
            })
    
    return results