Spaces:
Sleeping
Sleeping
File size: 4,417 Bytes
cbe90ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | using System;
using System.Linq;
using System.Threading.Tasks;
using Microsoft.AspNetCore.Mvc;
using MongoDB.Driver;
using TelegramSaaS.Domain.Entities;
using TelegramSaaS.Infrastructure.Persistence;
namespace TelegramSaaS.API.Controllers;
[ApiController]
[Route("api/v1/gpu")]
public class GpuRegistrationController : ControllerBase
{
private readonly MongoDbContext _dbContext;
private const string CollectionName = "GpuServers";
private static readonly TimeSpan HeartbeatTimeout = TimeSpan.FromHours(6);
public GpuRegistrationController(MongoDbContext dbContext)
{
_dbContext = dbContext;
}
[HttpPost("register")]
public async Task<IActionResult> Register([FromBody] GpuRegisterRequest request)
{
if (string.IsNullOrWhiteSpace(request.Url) || string.IsNullOrWhiteSpace(request.Label))
return BadRequest(new { error = "Both 'url' and 'label' are required." });
var label = request.Label.Trim().ToLowerInvariant();
var collection = _dbContext.GetCollection<GpuServer>(CollectionName);
var filter = Builders<GpuServer>.Filter.Eq(s => s.Label, label);
var existing = await collection.Find(filter).FirstOrDefaultAsync();
if (existing != null)
{
var update = Builders<GpuServer>.Update
.Set(s => s.Url, request.Url.Trim())
.Set(s => s.LastHeartbeat, DateTime.UtcNow)
.Set(s => s.IsActive, true)
.Set(s => s.UpdatedAt, DateTime.UtcNow);
await collection.UpdateOneAsync(filter, update);
}
else
{
var server = new GpuServer
{
Url = request.Url.Trim(),
Label = label,
LastHeartbeat = DateTime.UtcNow,
IsActive = true
};
await collection.InsertOneAsync(server);
}
return Ok(new { message = $"GPU server '{label}' registered successfully.", url = request.Url.Trim() });
}
[HttpPost("heartbeat")]
public async Task<IActionResult> Heartbeat([FromBody] GpuHeartbeatRequest request)
{
if (string.IsNullOrWhiteSpace(request.Url))
return BadRequest(new { error = "'url' is required." });
var collection = _dbContext.GetCollection<GpuServer>(CollectionName);
var filter = Builders<GpuServer>.Filter.Eq(s => s.Url, request.Url.Trim());
var update = Builders<GpuServer>.Update
.Set(s => s.LastHeartbeat, DateTime.UtcNow)
.Set(s => s.IsActive, true)
.Set(s => s.UpdatedAt, DateTime.UtcNow);
var result = await collection.UpdateOneAsync(filter, update);
if (result.MatchedCount == 0)
return NotFound(new { error = "GPU server not found." });
return Ok(new { message = "Heartbeat updated." });
}
[HttpGet("urls")]
public async Task<IActionResult> GetActiveUrls()
{
var collection = _dbContext.GetCollection<GpuServer>(CollectionName);
var cutoff = DateTime.UtcNow.Add(-HeartbeatTimeout);
var filter = Builders<GpuServer>.Filter.And(
Builders<GpuServer>.Filter.Eq(s => s.IsActive, true),
Builders<GpuServer>.Filter.Gt(s => s.LastHeartbeat, cutoff)
);
var servers = await collection.Find(filter).ToListAsync();
var result = servers.Select(s => new
{
url = s.Url,
label = s.Label,
lastHeartbeat = s.LastHeartbeat
});
return Ok(new { servers = result });
}
[HttpDelete("unregister")]
public async Task<IActionResult> Unregister([FromBody] GpuUnregisterRequest request)
{
if (string.IsNullOrWhiteSpace(request.Url))
return BadRequest(new { error = "'url' is required." });
var collection = _dbContext.GetCollection<GpuServer>(CollectionName);
var filter = Builders<GpuServer>.Filter.Eq(s => s.Url, request.Url.Trim());
var result = await collection.DeleteOneAsync(filter);
if (result.DeletedCount == 0)
return NotFound(new { error = "GPU server not found." });
return Ok(new { message = "GPU server unregistered." });
}
}
// Request DTOs
public record GpuRegisterRequest(string Url, string Label);
public record GpuHeartbeatRequest(string Url);
public record GpuUnregisterRequest(string Url);
|