mirror of
https://github.com/hanzoai/visor.git
synced 2026-09-24 17:04:14 +00:00
Aishu KamalandHanzo Dev
4cbdcc589b
Add a CLI tool and Go library that speak the libtpu checkpoint/restore control protocol (pkg/sentry/control/tpu_control.proto) against unsandboxed processes. gVisor coordinates libtpu C/R during sandbox save/restore (pkg/sentry/control/tpu_control.go), but there is no way to exercise or debug that protocol path without a full runsc save cycle. This tool fills the gap: it discovers libtpu control threads by scanning /proc/<pid>/task/*/comm, opens the encoded pipe FDs via /proc/<pid>/fd/, and sends ControlRequest messages. Accepts a comma-separated PID list and runs actions concurrently across PIDs with per-PID results, for multi-process jobs that need all processes checkpointed together. Co-authored-by: Hanzo Dev <[email protected]>
89 lines
2.6 KiB
Go
89 lines
2.6 KiB
Go
// Copyright 2026 The gVisor Authors.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package tpucontrol
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
|
|
tpupb "gvisor.dev/gvisor/pkg/sentry/control/tpu_control_proto_go_proto"
|
|
)
|
|
|
|
// actionName returns a short lower-case name for a control action, e.g.
|
|
// "checkpoint" for ACTION_CHECKPOINT.
|
|
func actionName(a tpupb.ControlAction) string {
|
|
return strings.ToLower(strings.TrimPrefix(a.String(), "ACTION_"))
|
|
}
|
|
|
|
// DiscoveryError indicates failure to find libtpu control threads.
|
|
type DiscoveryError struct {
|
|
PID int
|
|
Err error
|
|
}
|
|
|
|
// Error implements error.Error.
|
|
func (e *DiscoveryError) Error() string {
|
|
return fmt.Sprintf("discovery for PID %d: %v", e.PID, e.Err)
|
|
}
|
|
|
|
// Unwrap supports errors.Is/As.
|
|
func (e *DiscoveryError) Unwrap() error { return e.Err }
|
|
|
|
// TimeoutError indicates the operation timed out waiting for a response.
|
|
type TimeoutError struct {
|
|
Action tpupb.ControlAction
|
|
TID int
|
|
Err error
|
|
}
|
|
|
|
// Error implements error.Error.
|
|
func (e *TimeoutError) Error() string {
|
|
return fmt.Sprintf("%s timed out on thread %d: %v", actionName(e.Action), e.TID, e.Err)
|
|
}
|
|
|
|
// Unwrap supports errors.Is/As.
|
|
func (e *TimeoutError) Unwrap() error { return e.Err }
|
|
|
|
// ProtocolError indicates a wire-format or response parsing failure.
|
|
type ProtocolError struct {
|
|
Action tpupb.ControlAction
|
|
TID int
|
|
Err error
|
|
}
|
|
|
|
// Error implements error.Error.
|
|
func (e *ProtocolError) Error() string {
|
|
return fmt.Sprintf("%s protocol error on thread %d: %v", actionName(e.Action), e.TID, e.Err)
|
|
}
|
|
|
|
// Unwrap supports errors.Is/As.
|
|
func (e *ProtocolError) Unwrap() error { return e.Err }
|
|
|
|
// ControlFailedError indicates libtpu returned success=false.
|
|
type ControlFailedError struct {
|
|
Action tpupb.ControlAction
|
|
TID int
|
|
CurrentState tpupb.RuntimeState
|
|
ErrorMessage string
|
|
}
|
|
|
|
// Error implements error.Error.
|
|
func (e *ControlFailedError) Error() string {
|
|
if e.ErrorMessage != "" {
|
|
return fmt.Sprintf("%s failed on thread %d: %s", actionName(e.Action), e.TID, e.ErrorMessage)
|
|
}
|
|
return fmt.Sprintf("%s failed on thread %d: state=%s", actionName(e.Action), e.TID, e.CurrentState)
|
|
}
|