// Package model provides TUI data structures and state management package model import ( "fmt" "time" "github.com/charmbracelet/bubbles/key" "github.com/charmbracelet/bubbles/list" "github.com/charmbracelet/bubbles/spinner" "github.com/charmbracelet/bubbles/textinput" "github.com/charmbracelet/bubbles/viewport" "github.com/charmbracelet/lipgloss" ) // ViewMode represents the current view mode in the TUI type ViewMode int // ViewMode constants represent different TUI views const ( ViewModeJobs ViewMode = iota // Jobs view mode ViewModeGPU // GPU status view mode ViewModeQueue // Queue status view mode ViewModeContainer // Container status view mode ViewModeSettings // Settings view mode ViewModeDatasets // Datasets view mode ViewModeExperiments // Experiments view mode ) // JobStatus represents the status of a job type JobStatus string // JobStatus constants represent different job states const ( StatusPending JobStatus = "pending" // Job is pending StatusQueued JobStatus = "queued" // Job is queued StatusRunning JobStatus = "running" // Job is running StatusFinished JobStatus = "finished" // Job is finished StatusFailed JobStatus = "failed" // Job is failed ) // Job represents a job in the TUI type Job struct { Name string Status JobStatus TaskID string Priority int64 } // Title returns the job title for display func (j Job) Title() string { return j.Name } // Description returns a formatted description with status icon func (j Job) Description() string { icon := map[JobStatus]string{ StatusPending: "⏸", StatusQueued: "⏳", StatusRunning: "▶", StatusFinished: "✓", StatusFailed: "✗", }[j.Status] pri := "" if j.Priority > 0 { pri = fmt.Sprintf(" [P%d]", j.Priority) } return fmt.Sprintf("%s %s%s", icon, j.Status, pri) } // FilterValue returns the value used for filtering func (j Job) FilterValue() string { return j.Name } // Task represents a task in the TUI type Task struct { ID string `json:"id"` JobName string `json:"job_name"` Args string `json:"args"` Status string `json:"status"` Priority int64 `json:"priority"` CreatedAt time.Time `json:"created_at"` StartedAt *time.Time `json:"started_at,omitempty"` EndedAt *time.Time `json:"ended_at,omitempty"` Error string `json:"error,omitempty"` Metadata map[string]string `json:"metadata,omitempty"` Tracking *TrackingConfig `json:"tracking,omitempty"` } // TrackingConfig specifies experiment tracking tools type TrackingConfig struct { MLflow *MLflowTrackingConfig `json:"mlflow,omitempty"` TensorBoard *TensorBoardTrackingConfig `json:"tensorboard,omitempty"` Wandb *WandbTrackingConfig `json:"wandb,omitempty"` } // MLflowTrackingConfig controls MLflow integration type MLflowTrackingConfig struct { Enabled bool `json:"enabled"` Mode string `json:"mode,omitempty"` TrackingURI string `json:"tracking_uri,omitempty"` } // TensorBoardTrackingConfig controls TensorBoard integration type TensorBoardTrackingConfig struct { Enabled bool `json:"enabled"` Mode string `json:"mode,omitempty"` } // WandbTrackingConfig controls Weights & Biases integration type WandbTrackingConfig struct { Enabled bool `json:"enabled"` Mode string `json:"mode,omitempty"` APIKey string `json:"api_key,omitempty"` Project string `json:"project,omitempty"` Entity string `json:"entity,omitempty"` } // DatasetInfo represents dataset information in the TUI type DatasetInfo struct { Name string `json:"name"` SizeBytes int64 `json:"size_bytes"` Location string `json:"location"` LastAccess time.Time `json:"last_access"` } // State holds the application state type State struct { Jobs []Job QueuedTasks []*Task Datasets []DatasetInfo JobList list.Model GpuView viewport.Model ContainerView viewport.Model QueueView viewport.Model SettingsView viewport.Model DatasetView viewport.Model ExperimentsView viewport.Model Input textinput.Model APIKeyInput textinput.Model Status string ErrorMsg string InputMode bool Width int Height int ShowHelp bool Spinner spinner.Model ActiveView ViewMode LastRefresh time.Time IsLoading bool JobStats map[JobStatus]int APIKey string SettingsIndex int Keys KeyMap } // KeyMap defines key bindings for the TUI type KeyMap struct { Refresh key.Binding Trigger key.Binding TriggerArgs key.Binding ViewQueue key.Binding ViewContainer key.Binding ViewGPU key.Binding ViewJobs key.Binding ViewDatasets key.Binding ViewExperiments key.Binding ViewSettings key.Binding Cancel key.Binding Delete key.Binding MarkFailed key.Binding RefreshGPU key.Binding Help key.Binding Quit key.Binding } // Keys contains the default key bindings for the TUI var Keys = KeyMap{ Refresh: key.NewBinding(key.WithKeys("r"), key.WithHelp("r", "refresh all")), Trigger: key.NewBinding(key.WithKeys("t"), key.WithHelp("t", "queue job")), TriggerArgs: key.NewBinding(key.WithKeys("a"), key.WithHelp("a", "queue w/ args")), ViewQueue: key.NewBinding(key.WithKeys("v"), key.WithHelp("v", "view queue")), ViewContainer: key.NewBinding(key.WithKeys("o"), key.WithHelp("o", "containers")), ViewGPU: key.NewBinding(key.WithKeys("g"), key.WithHelp("g", "gpu status")), ViewJobs: key.NewBinding(key.WithKeys("1"), key.WithHelp("1", "job list")), ViewDatasets: key.NewBinding(key.WithKeys("2"), key.WithHelp("2", "datasets")), ViewExperiments: key.NewBinding(key.WithKeys("3"), key.WithHelp("3", "experiments")), Cancel: key.NewBinding(key.WithKeys("c"), key.WithHelp("c", "cancel task")), Delete: key.NewBinding(key.WithKeys("d"), key.WithHelp("d", "delete job")), MarkFailed: key.NewBinding(key.WithKeys("f"), key.WithHelp("f", "mark failed")), RefreshGPU: key.NewBinding(key.WithKeys("G"), key.WithHelp("G", "refresh GPU")), ViewSettings: key.NewBinding(key.WithKeys("s"), key.WithHelp("s", "settings")), Help: key.NewBinding(key.WithKeys("h", "?"), key.WithHelp("h/?", "toggle help")), Quit: key.NewBinding(key.WithKeys("q", "ctrl+c"), key.WithHelp("q", "quit")), } // InitialState creates the initial application state func InitialState(apiKey string) State { items := []list.Item{} delegate := list.NewDefaultDelegate() delegate.Styles.SelectedTitle = delegate.Styles.SelectedTitle. Foreground(lipgloss.Color("170")). Bold(true) delegate.Styles.SelectedDesc = delegate.Styles.SelectedDesc. Foreground(lipgloss.Color("246")) jobList := list.New(items, delegate, 0, 0) jobList.Title = "ML Jobs & Queue" jobList.SetShowStatusBar(true) jobList.SetFilteringEnabled(true) jobList.SetShowHelp(false) // Styles will be set in View or here? // Keeping style initialization here as it's part of the model state setup jobList.Styles.Title = lipgloss.NewStyle(). Bold(true). Foreground(lipgloss.AdaptiveColor{Light: "#2980b9", Dark: "#7aa2f7"}). Padding(0, 0, 1, 0) input := textinput.New() input.Placeholder = "Args: --epochs 100 --lr 0.001 --priority 5" input.Width = 60 input.CharLimit = 200 apiKeyInput := textinput.New() apiKeyInput.Placeholder = "Enter API key..." apiKeyInput.Width = 40 apiKeyInput.CharLimit = 200 s := spinner.New() s.Spinner = spinner.Dot s.Style = lipgloss.NewStyle().Foreground(lipgloss.AdaptiveColor{Light: "#2980b9", Dark: "#7aa2f7"}) return State{ JobList: jobList, GpuView: viewport.New(0, 0), ContainerView: viewport.New(0, 0), QueueView: viewport.New(0, 0), SettingsView: viewport.New(0, 0), DatasetView: viewport.New(0, 0), ExperimentsView: viewport.New(0, 0), Input: input, APIKeyInput: apiKeyInput, Status: "Connected", InputMode: false, ShowHelp: false, Spinner: s, ActiveView: ViewModeJobs, LastRefresh: time.Now(), IsLoading: false, JobStats: make(map[JobStatus]int), APIKey: apiKey, SettingsIndex: 0, Keys: Keys, } }