Help us shape the future of Go! Take the 2026 Go Developer Survey

client

package
v0.10.1 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Oct 23, 2026 License: MIT Imports: 38 Imported by: 0

Documentation

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func DiscoverGPUCount added in v0.8.9

func DiscoverGPUCount() int

DiscoverGPUCount returns the number of NVIDIA GPUs visible on the host (counted via `nvidia-smi -L`). Returns 0 when nvidia-smi is missing or fails — the typical state on CPU-only flavors. Logs the outcome at INFO level so it lands in worker stdout for ops grep.

Why nvidia-smi instead of reading /dev/nvidia*: nvidia-smi works regardless of how the device files are named (some images use /dev/nvidiactl + /dev/nvidia-uvm without per-GPU node files), and the line count IS the gospel definition Azure / the kernel agrees on. Cheap (< 100 ms) and only runs once at startup.

func ParseCpusetMems added in v0.8.5

func ParseCpusetMems(s string) ([]int, error)

ParseCpusetMems parses docker's --cpuset-mems output ("0,2,3") back into a slice of node IDs. Used by restart recovery.

func Run

func Run(ctx context.Context, serverAddr string, concurrency int32, name, store, token string, isPermanent bool, provider *string, region *string, maxCPU int32, maxMem float32, noBare ...bool) error

/ client launcher

func VisibleDevicesEnv added in v0.8.9

func VisibleDevicesEnv(devs []int) string

VisibleDevicesEnv formats a slice of device indices as the value for CUDA_VISIBLE_DEVICES (comma-separated, ascending). Returns the empty string for an empty slice.

Types

type Allocation added in v0.8.5

type Allocation struct {
	Cpuset   string // ready for --cpuset-cpus
	Memset   string // ready for --cpuset-mems
	CPUCount int    // number of CPUs in Cpuset
	Nodes    []int  // node IDs (for logging / debugging)
}

Allocation describes what NumaAllocator.Allocate handed out. Fields map directly onto docker flags (Cpuset → --cpuset-cpus, Memset → --cpuset-mems) plus the cpu count so the executor can override the per-task $CPU env var to match.

type DownloadFailure

type DownloadFailure struct {
	Task    *pb.Task
	Message string
}

DownloadFailure carries a unified failure notification from the downloader to a single status update point in the client loop.

type DownloadManager

type DownloadManager struct {
	ResourceDownloads map[string][]*pb.Task   // SourcePath → Tasks waiting
	ResourceMemory    map[string]FileMetadata // SourcePath → Metadata

	FileQueue       chan *FileTransfer    // Limited queue (maxDownloads=20)
	TaskQueue       chan *pb.Task         // Unlimited queue (tasks waiting for download)
	CompletionQueue chan *FileMetadata    // Completed downloads
	ExecQueue       chan *pb.Task         // Tasks ready for execution
	FailedQueue     chan *DownloadFailure // Tasks that failed during downloads (centralized reporting)
	Store           string

	// Tracks tasks currently being scheduled for downloads to ensure
	// idempotency at the boundary. All WRITES happen inside the
	// DownloadManager's single dispatch loop (handleNewTask /
	// handleFileCompletion), but client.workerLoop reads it from
	// another goroutine to drop server-side duplicates — a plain
	// map there panics with "concurrent map read and map write".
	// sync.Map is the right shape: write-mostly from one goroutine,
	// cheap presence-check Load from another. Value is struct{} since
	// only the key's presence matters.
	EnqueuedTasks sync.Map
	// Per-task chrono map for download durations
	DownloadStart map[int32]time.Time
	// Per-task pending set for download items
	TaskPending   map[int32]map[string]bool // Task ID → set of pending item keys (I:input, R:resource, D:docker)
	RcloneRemotes *pb.RcloneRemotes
	// contains filtered or unexported fields
}

DownloadManager manages task downloads.

func NewDownloadManager

func NewDownloadManager(store string, reporter *event.Reporter, rcloneRemotes *pb.RcloneRemotes) *DownloadManager

NewDownloadManager initializes the download manager.

func RunDownloader

func RunDownloader(store string, reporter *event.Reporter, rcloneRemotes *pb.RcloneRemotes) *DownloadManager

run downloader

func (*DownloadManager) ProcessDownloads

func (dm *DownloadManager) ProcessDownloads()

ProcessDownloads manages incoming tasks and processes downloads.

func (*DownloadManager) StartDownloadWorkers

func (dm *DownloadManager) StartDownloadWorkers()

StartDownloadWorkers launches download workers.

type FileMetadata

type FileMetadata struct {
	TaskId     int32
	Task       *pb.Task
	SourcePath string
	FilePath   string
	FileType   FileType
	Size       int64
	Date       time.Time
	MD5        string
	Success    bool // true if download succeeded
}

ResourceMetadata stores metadata for downloaded resources.

type FileTransfer

type FileTransfer struct {
	TaskId     int32
	Task       *pb.Task
	FileType   FileType
	SourcePath string
	TargetPath string
}

type FileType

type FileType int

FileTransfer represents a download task.

const (
	InputFile FileType = iota
	ResourceFile
	DockerImage
)

type GPUAllocator added in v0.8.9

type GPUAllocator struct {
	// contains filtered or unexported fields
}

GPUAllocator partitions a worker's GPU devices into per-task slots driven by task_spec.gpu. Mirrors NumaAllocator's shape: in-memory only, concurrency-safe, slots keyed by device index 0..N-1. A worker restart re-zeros the map — same risk profile as the NUMA allocator (a crashed worker could orphan device assignments until the next restart). For typical scitq deploys this is fine because the worker is bounced when the binary or the host driver changes.

The allocator's contract:

Allocate(taskID, n) reserves N free devices for taskID and
  returns their indices. Returns (nil, false) if there aren't N
  free devices, or if the worker has no GPUs at all. Caller is
  responsible for calling Release(taskID) when the task ends.
Release(taskID) frees every device that was reserved for taskID.
  Idempotent — safe to call on tasks that never got an allocation
  (e.g. CPU-only tasks).
DeviceCount() returns the worker's total GPU count (the size of
  the slot map). 0 on CPU-only workers — callers use it to gate
  whether to attempt allocation at all.

func NewGPUAllocator added in v0.8.9

func NewGPUAllocator(count int) *GPUAllocator

NewGPUAllocator builds an allocator for a worker with `count` GPU devices (0..count-1). count=0 yields a no-op allocator that always returns (nil, false) on Allocate — used on CPU-only workers so callers can hold a non-nil reference unconditionally.

func (*GPUAllocator) Allocate added in v0.8.9

func (a *GPUAllocator) Allocate(taskID int32, n int) ([]int, bool)

Allocate finds `n` free GPU devices and binds them to `taskID`. The returned slice is the list of device indices (ascending) the task owns. Caller MUST call Release(taskID) on completion to free them. Returns (nil, false) if n exceeds the worker's GPU count, no free slots, or n < 1.

func (*GPUAllocator) DeviceCount added in v0.8.9

func (a *GPUAllocator) DeviceCount() int

DeviceCount returns the worker's total GPU count.

func (*GPUAllocator) FreeCount added in v0.8.9

func (a *GPUAllocator) FreeCount() int

FreeCount returns the number of currently-free GPU devices. Used by the workerLoop's fetch-side capacity check to avoid pulling more GPU-needing tasks than the worker can host.

func (*GPUAllocator) Release added in v0.8.9

func (a *GPUAllocator) Release(taskID int32)

Release frees every device currently bound to `taskID`. Safe to call on a task that never allocated (e.g. CPU-only tasks).

func (*GPUAllocator) String added in v0.8.9

func (a *GPUAllocator) String() string

String summarises the allocator state for logging.

type LiveCaps added in v0.8.5

type LiveCaps struct {
	// contains filtered or unexported fields
}

LiveCaps holds the worker's resource caps in a form that can be read by per-task goroutines (excuterThread) and updated from the fetchTasks loop on every ping. Atomic so neither side needs to take a lock on the hot path.

Zero (the default) means "autodetect": fall back to runtime.NumCPU() / host memory at read time, matching the static-flag behaviour.

func NewLiveCaps added in v0.8.5

func NewLiveCaps(cpu int32, memGiB float32) *LiveCaps

func (*LiveCaps) CPU added in v0.8.5

func (l *LiveCaps) CPU() int32

func (*LiveCaps) Mem added in v0.8.5

func (l *LiveCaps) Mem() float32

func (*LiveCaps) SetCPU added in v0.8.5

func (l *LiveCaps) SetCPU(v int32)

func (*LiveCaps) SetMem added in v0.8.5

func (l *LiveCaps) SetMem(v float32)

type NumaAllocator added in v0.8.5

type NumaAllocator struct {
	// contains filtered or unexported fields
}

NumaAllocator tracks which NUMA node currently hosts which task. The map is purely in-memory: a worker restart rebuilds it from `docker inspect` of still-running scitq containers (see RebuildFromDocker). All methods are concurrency-safe — both the fetcher goroutine and the executer goroutine touch this.

func NewNumaAllocator added in v0.8.5

func NewNumaAllocator(topo *NumaTopology) *NumaAllocator

func (*NumaAllocator) Allocate added in v0.8.5

func (a *NumaAllocator) Allocate(taskID int32, n int) (Allocation, bool)

Allocate finds `n` free NUMA nodes and binds them to `taskID`. Returns the docker --cpuset-cpus / --cpuset-mems values plus the CPU count, or (zero, false) if allocation isn't possible (n exceeds the topology, no free slots, or the host has no NUMA topology at all). Caller is responsible for releasing the slots via Release on task completion.

func (*NumaAllocator) EffectiveNodeCount added in v0.8.5

func (a *NumaAllocator) EffectiveNodeCount() int

EffectiveNodeCount is NumaNodeCount() with a floor of 1. For the purpose of *deriving concurrency* from task_spec.numa, a non-NUMA host is treated as if it had a single NUMA node: a step with `numa: 1` runs one task at a time on it (and falls through the allocator's unbound-fallback at dispatch time). This keeps the DSL usable across heterogeneous fleets — a YAML written for EPYC still works (with reduced parallelism) on a single-socket dev box.

func (*NumaAllocator) MarkAllocated added in v0.8.5

func (a *NumaAllocator) MarkAllocated(taskID int32, nodeIDs []int)

MarkAllocated forces nodes into the bound state. Used by restart recovery: when the worker comes back up, it consults docker for containers still running and records their NUMA assignment so the allocator doesn't hand the same slots out again.

func (*NumaAllocator) NumaNodeCount added in v0.8.5

func (a *NumaAllocator) NumaNodeCount() int

NumaNodeCount returns the number of NUMA nodes the worker can pin to, or 0 if the host is NUMA-unaware. Callers use this to gate physical binding (allocator returns false when count is 0).

func (*NumaAllocator) RebuildFromDocker added in v0.8.5

func (a *NumaAllocator) RebuildFromDocker(ctx context.Context, workerName string)

RebuildFromDocker scans this worker's still-running scitq containers and re-claims any NUMA slots they hold. Used on worker startup so a process restart in the middle of a task doesn't release node IDs that are still genuinely in use by a running container.

Container name format (set in executeTask): "scitq-<workerName>-task-<id>". We read HostConfig.CpusetMems and translate it back into node IDs.

func (*NumaAllocator) Release added in v0.8.5

func (a *NumaAllocator) Release(taskID int32)

Release frees every NUMA node currently bound to taskID. No-op if the task wasn't bound (host had no topology, or allocation failed and we ran unbound).

type NumaNode added in v0.8.5

type NumaNode struct {
	ID      int
	CPUList string // ready for --cpuset-cpus
}

NumaNode describes one NUMA node on the host. CPUList is kept as the raw cpulist string from /sys (e.g. "0-5,12-17") so we can feed it straight into docker `--cpuset-cpus` without parsing.

type NumaTopology added in v0.8.5

type NumaTopology struct {
	Nodes []NumaNode
}

NumaTopology is the host's NUMA layout. nil = NUMA-unaware host (sysfs reported one or zero nodes, or wasn't readable). Treat a nil topology as "the worker cannot honour task.numa" — pinning becomes a no-op with a single warning logged at dispatch time.

func DiscoverNumaTopology added in v0.8.5

func DiscoverNumaTopology() *NumaTopology

DiscoverNumaTopology reads /sys/devices/system/node/node*/cpulist once at startup. Returns nil (not an error) on a host with no NUMA information — the worker can still run, just without binding.

type SyncCounter

type SyncCounter struct {
	// contains filtered or unexported fields
}

func (*SyncCounter) Decrement

func (sc *SyncCounter) Decrement(taskID int32) (int, bool)

func (*SyncCounter) Set

func (sc *SyncCounter) Set(taskID int32, count int)

type UploadFailure

type UploadFailure struct {
	Task    *pb.Task
	Message string
}

UploadFailure carries a unified failure notification from the uploader to a single status update point in the client loop.

type UploadFile

type UploadFile struct {
	Task       *pb.Task
	SourcePath string
	TargetPath string
}

type UploadManager

type UploadManager struct {
	UploadQueue chan *UploadFile
	Completion  chan *pb.Task
	FailedQueue chan *UploadFailure
	Store       string
	Client      pb.TaskQueueClient

	RcloneRemotes *pb.RcloneRemotes
	// contains filtered or unexported fields
}

func NewUploadManager

func NewUploadManager(store string, client pb.TaskQueueClient, reporter *event.Reporter, rcloneRemotes *pb.RcloneRemotes) *UploadManager

func RunUploader

func RunUploader(store string, client pb.TaskQueueClient, activeTasks *sync.Map, reporter *event.Reporter, rcloneRemotes *pb.RcloneRemotes) *UploadManager

RunUploader starts upload workers and returns the manager.

func (*UploadManager) EnqueueTaskOutput

func (um *UploadManager) EnqueueTaskOutput(task *pb.Task)

func (*UploadManager) StartUploadWorkers

func (um *UploadManager) StartUploadWorkers(activeTasks *sync.Map)

type WorkerConfig

type WorkerConfig struct {
	WorkerId    int32
	ServerAddr  string
	Concurrency int32
	Name        string
	Store       string
	Token       string
	IsPermanent bool
	Provider    *string
	Region      *string
	NoBare      bool // reject bare tasks
	// MaxCPU / MaxMem cap what this worker advertises to scitq, for
	// shared nodes that should only contribute part of their hardware.
	// Zero means "autodetect" (runtime.NumCPU / total memory). Validated
	// against the real machine in registerSpecs — a value larger than
	// the hardware is rejected, since the server uses these numbers
	// directly for flavor matching and $CPU/$MEM injection.
	MaxCPU int32
	MaxMem float32 // GiB
}

WorkerConfig holds worker settings from CLI args.

Directories

Path Synopsis
Package iothrottle implements adaptive per-worker concurrency throttling based on continuously-sampled iowait%.
Package iothrottle implements adaptive per-worker concurrency throttling based on continuously-sampled iowait%.
Package peakmem reads the kernel-tracked peak resident memory of a finished task.
Package peakmem reads the kernel-tracked peak resident memory of a finished task.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL