toki
3624ff0f01
- Add inference provider extension phase and milestones for node multi-target serving - Update edge cmd config and node mapper for multi-target support - Add Ollama adapter updates and transport client changes - Update config package and edge.yaml for new settings - Extend runtime proto with multi-target serving fields - Add agent task archives and plans for adapter registry, route catalog, engine profiles, edge options, commands/caps, runtime concurrency, and field smoke tests
116 lines
3.1 KiB
Go
116 lines
3.1 KiB
Go
package transport
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net"
|
|
"strconv"
|
|
"time"
|
|
|
|
toki "git.toki-labs.com/toki/proto-socket/go"
|
|
"go.uber.org/zap"
|
|
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
const (
|
|
heartbeatIntervalSec = 30
|
|
// heartbeatWaitSec is kept above heartbeatIntervalSec as defence in depth:
|
|
// the library wait-timer callback already self-heals stale state (see
|
|
// proto-socket go/base_client.go sendHeartBeat), but a larger wait window
|
|
// gives the peer's next heartbeat an extra chance to overwrite any stray
|
|
// timer on slow or jittery links before it fires.
|
|
heartbeatWaitSec = 45
|
|
registerTimeout = 10 * time.Second
|
|
registerInitialWait = 100 * time.Millisecond
|
|
registerRetryWait = 250 * time.Millisecond
|
|
registerAttempts = 3
|
|
)
|
|
|
|
// RegisterResult is returned by DialEdge after successful registration.
|
|
type RegisterResult struct {
|
|
Session *Session
|
|
NodeID string
|
|
Alias string
|
|
Config *iop.NodeConfigPayload
|
|
}
|
|
|
|
// DialEdge connects to edge, performs the registration handshake, and returns
|
|
// a RegisterResult. Call result.Session.SetHandler after creating node.Node.
|
|
func DialEdge(ctx context.Context, addr, token string, logger *zap.Logger) (*RegisterResult, error) {
|
|
host, portStr, err := net.SplitHostPort(addr)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("transport: invalid addr %q: %w", addr, err)
|
|
}
|
|
port, err := strconv.Atoi(portStr)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("transport: invalid port %q: %w", portStr, err)
|
|
}
|
|
|
|
client, err := toki.DialTcp(ctx, host, port, heartbeatIntervalSec, heartbeatWaitSec, nodeParserMap())
|
|
if err != nil {
|
|
return nil, fmt.Errorf("transport: dial edge %s: %w", addr, err)
|
|
}
|
|
|
|
resp, err := registerWithEdge(ctx, client, token, logger)
|
|
if err != nil {
|
|
_ = client.Close()
|
|
return nil, fmt.Errorf("transport: register: %w", err)
|
|
}
|
|
if !resp.GetAccepted() {
|
|
_ = client.Close()
|
|
return nil, fmt.Errorf("transport: register rejected: %s", resp.GetReason())
|
|
}
|
|
|
|
sess := newSession(client, logger, resp.GetNodeId(), resp.GetAlias())
|
|
logger.Info("registered with edge",
|
|
zap.String("node_id", resp.GetNodeId()),
|
|
zap.String("alias", resp.GetAlias()),
|
|
)
|
|
return &RegisterResult{
|
|
Session: sess,
|
|
NodeID: resp.GetNodeId(),
|
|
Alias: resp.GetAlias(),
|
|
Config: resp.GetConfig(),
|
|
}, nil
|
|
}
|
|
|
|
func registerWithEdge(ctx context.Context, client *toki.TcpClient, token string, logger *zap.Logger) (*iop.RegisterResponse, error) {
|
|
timer := time.NewTimer(registerInitialWait)
|
|
select {
|
|
case <-ctx.Done():
|
|
timer.Stop()
|
|
return nil, ctx.Err()
|
|
case <-timer.C:
|
|
}
|
|
|
|
var lastErr error
|
|
for attempt := 1; attempt <= registerAttempts; attempt++ {
|
|
resp, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
|
|
&client.Communicator,
|
|
&iop.RegisterRequest{Token: token},
|
|
registerTimeout,
|
|
)
|
|
if err == nil {
|
|
return resp, nil
|
|
}
|
|
lastErr = err
|
|
if attempt == registerAttempts || !client.IsAlive() {
|
|
break
|
|
}
|
|
if logger != nil {
|
|
logger.Warn("register request failed, retrying",
|
|
zap.Int("attempt", attempt),
|
|
zap.Error(err),
|
|
)
|
|
}
|
|
timer := time.NewTimer(registerRetryWait)
|
|
select {
|
|
case <-ctx.Done():
|
|
timer.Stop()
|
|
return nil, ctx.Err()
|
|
case <-timer.C:
|
|
}
|
|
}
|
|
return nil, lastErr
|
|
}
|