package jupyter import ( "fmt" "os" "path/filepath" "strings" "time" "github.com/jfraeys/fetch_ml/internal/logging" ) const ( mountTypeBind = "bind" ) // WorkspaceManager handles workspace mounting and volume management for Jupyter services type WorkspaceManager struct { logger *logging.Logger basePath string mounts map[string]*WorkspaceMount } // WorkspaceMount represents a workspace mount configuration type WorkspaceMount struct { ID string `json:"id"` HostPath string `json:"host_path"` ContainerPath string `json:"container_path"` MountType string `json:"mount_type"` // "bind", "volume", "tmpfs" ReadOnly bool `json:"read_only"` Options map[string]string `json:"options"` Services []string `json:"services"` // Service IDs using this mount } // MountRequest defines parameters for creating a workspace mount type MountRequest struct { HostPath string `json:"host_path"` ContainerPath string `json:"container_path"` MountType string `json:"mount_type"` ReadOnly bool `json:"read_only"` Options map[string]string `json:"options"` } // NewWorkspaceManager creates a new workspace manager func NewWorkspaceManager(logger *logging.Logger, basePath string) *WorkspaceManager { return &WorkspaceManager{ logger: logger, basePath: basePath, mounts: make(map[string]*WorkspaceMount), } } // CreateWorkspaceMount creates a new workspace mount func (wm *WorkspaceManager) CreateWorkspaceMount(req *MountRequest) (*WorkspaceMount, error) { // Validate request if err := wm.validateMountRequest(req); err != nil { return nil, err } // Generate mount ID mountID := wm.generateMountID(req.HostPath) // Ensure host path exists if req.MountType == mountTypeBind { if err := wm.ensureHostPath(req.HostPath); err != nil { return nil, fmt.Errorf("failed to prepare host path: %w", err) } } // Create mount record mount := &WorkspaceMount{ ID: mountID, HostPath: req.HostPath, ContainerPath: req.ContainerPath, MountType: req.MountType, ReadOnly: req.ReadOnly, Options: req.Options, Services: []string{}, } // Store mount wm.mounts[mountID] = mount wm.logger.Info("workspace mount created", "mount_id", mountID, "host_path", req.HostPath, "container_path", req.ContainerPath, "mount_type", req.MountType) return mount, nil } // GetMount retrieves a mount by ID func (wm *WorkspaceManager) GetMount(mountID string) (*WorkspaceMount, error) { mount, exists := wm.mounts[mountID] if !exists { return nil, fmt.Errorf("mount %s not found", mountID) } return mount, nil } // FindMountByPath finds a mount by host path func (wm *WorkspaceManager) FindMountByPath(hostPath string) (*WorkspaceMount, error) { for _, mount := range wm.mounts { if mount.HostPath == hostPath { return mount, nil } } return nil, fmt.Errorf("no mount found for path %s", hostPath) } // ListMounts returns all mounts func (wm *WorkspaceManager) ListMounts() []*WorkspaceMount { mounts := make([]*WorkspaceMount, 0, len(wm.mounts)) for _, mount := range wm.mounts { mounts = append(mounts, mount) } return mounts } // RemoveMount removes a workspace mount func (wm *WorkspaceManager) RemoveMount(mountID string) error { mount, exists := wm.mounts[mountID] if !exists { return fmt.Errorf("mount %s not found", mountID) } // Check if mount is in use if len(mount.Services) > 0 { return fmt.Errorf("cannot remove mount %s: in use by services %v", mountID, mount.Services) } // Remove mount delete(wm.mounts, mountID) wm.logger.Info("workspace mount removed", "mount_id", mountID, "host_path", mount.HostPath) return nil } // AttachService attaches a service to a mount func (wm *WorkspaceManager) AttachService(mountID, serviceID string) error { mount, exists := wm.mounts[mountID] if !exists { return fmt.Errorf("mount %s not found", mountID) } // Check if service is already attached for _, service := range mount.Services { if service == serviceID { return nil // Already attached } } // Attach service mount.Services = append(mount.Services, serviceID) wm.logger.Debug("service attached to mount", "mount_id", mountID, "service_id", serviceID) return nil } // DetachService detaches a service from a mount func (wm *WorkspaceManager) DetachService(mountID, serviceID string) error { mount, exists := wm.mounts[mountID] if !exists { return fmt.Errorf("mount %s not found", mountID) } // Find and remove service for i, service := range mount.Services { if service == serviceID { mount.Services = append(mount.Services[:i], mount.Services[i+1:]...) wm.logger.Debug("service detached from mount", "mount_id", mountID, "service_id", serviceID) return nil } } return nil // Service not attached } // GetMountsForService returns all mounts used by a service func (wm *WorkspaceManager) GetMountsForService(serviceID string) []*WorkspaceMount { var serviceMounts []*WorkspaceMount for _, mount := range wm.mounts { for _, service := range mount.Services { if service == serviceID { serviceMounts = append(serviceMounts, mount) break } } } return serviceMounts } // PrepareWorkspace prepares a workspace for Jupyter func (wm *WorkspaceManager) PrepareWorkspace(workspacePath string) (*WorkspaceMount, error) { // Check if mount already exists mount, err := wm.FindMountByPath(workspacePath) if err == nil { return mount, nil } // Create new mount req := &MountRequest{ HostPath: workspacePath, ContainerPath: "/workspace", MountType: mountTypeBind, ReadOnly: false, Options: map[string]string{ "Z": "", // For SELinux compatibility }, } return wm.CreateWorkspaceMount(req) } // validateMountRequest validates a mount request func (wm *WorkspaceManager) validateMountRequest(req *MountRequest) error { if req.HostPath == "" { return fmt.Errorf("host path is required") } if req.ContainerPath == "" { return fmt.Errorf("container path is required") } if req.MountType == "" { req.MountType = mountTypeBind } // Validate mount type validTypes := []string{"bind", "volume", "tmpfs"} valid := false for _, t := range validTypes { if req.MountType == t { valid = true break } } if !valid { return fmt.Errorf("invalid mount type %s, must be one of: %v", req.MountType, validTypes) } // For bind mounts, host path must be absolute if req.MountType == mountTypeBind && !filepath.IsAbs(req.HostPath) { return fmt.Errorf("bind mount host path must be absolute: %s", req.HostPath) } // Check for duplicate mounts for _, mount := range wm.mounts { if mount.HostPath == req.HostPath { return fmt.Errorf("mount for path %s already exists", req.HostPath) } } return nil } // generateMountID generates a unique mount ID func (wm *WorkspaceManager) generateMountID(hostPath string) string { // Create a safe ID from the host path safePath := strings.ToLower(hostPath) safePath = strings.ReplaceAll(safePath, "/", "-") safePath = strings.ReplaceAll(safePath, " ", "-") safePath = strings.ReplaceAll(safePath, "_", "-") // Remove leading dash safePath = strings.TrimPrefix(safePath, "-") return fmt.Sprintf("mount-%s", safePath) } // ensureHostPath ensures the host path exists and has proper permissions func (wm *WorkspaceManager) ensureHostPath(hostPath string) error { // Check if path exists info, err := os.Stat(hostPath) if err != nil { if os.IsNotExist(err) { // Create directory if err := os.MkdirAll(hostPath, 0750); err != nil { return fmt.Errorf("failed to create directory %s: %w", hostPath, err) } wm.logger.Info("created workspace directory", "path", hostPath) } else { return fmt.Errorf("failed to stat path %s: %w", hostPath, err) } } else if !info.IsDir() { return fmt.Errorf("host path %s is not a directory", hostPath) } // Check permissions if err := os.Chmod(hostPath, 0600); err != nil { wm.logger.Warn("failed to set permissions on workspace", "path", hostPath, "error", err) } return nil } // ValidateWorkspace validates a workspace for Jupyter use func (wm *WorkspaceManager) ValidateWorkspace(workspacePath string) error { // Check if workspace exists info, err := os.Stat(workspacePath) if err != nil { if os.IsNotExist(err) { return fmt.Errorf("workspace %s does not exist", workspacePath) } return fmt.Errorf("failed to access workspace %s: %w", workspacePath, err) } if !info.IsDir() { return fmt.Errorf("workspace %s is not a directory", workspacePath) } // Check for common Jupyter files jupyterFiles := []string{ "*.ipynb", "requirements.txt", "environment.yml", "Pipfile", "pyproject.toml", } var foundFiles []string for _, pattern := range jupyterFiles { matches, err := filepath.Glob(filepath.Join(workspacePath, pattern)) if err == nil && len(matches) > 0 { foundFiles = append(foundFiles, pattern) } } if len(foundFiles) == 0 { wm.logger.Warn("workspace may not be a Jupyter project", "path", workspacePath, "no_files_found", strings.Join(jupyterFiles, ", ")) } else { wm.logger.Info("workspace validated", "path", workspacePath, "found_files", strings.Join(foundFiles, ", ")) } return nil } // GetWorkspaceInfo returns information about a workspace func (wm *WorkspaceManager) GetWorkspaceInfo(workspacePath string) (*WorkspaceInfo, error) { info := &WorkspaceInfo{ Path: workspacePath, } // Get directory info dirInfo, err := os.Stat(workspacePath) if err != nil { return nil, fmt.Errorf("failed to stat workspace: %w", err) } info.Size = dirInfo.Size() info.Modified = dirInfo.ModTime() // Count files err = filepath.Walk(workspacePath, func(path string, fi os.FileInfo, err error) error { if err != nil { return err } if !fi.IsDir() { info.FileCount++ info.TotalSize += fi.Size() // Categorize files ext := strings.ToLower(filepath.Ext(path)) switch ext { case ".py": info.PythonFiles++ case ".ipynb": info.NotebookFiles++ case ".txt", ".md": info.TextFiles++ case ".json", ".yaml", ".yml": info.ConfigFiles++ } } return nil }) if err != nil { return nil, fmt.Errorf("failed to scan workspace: %w", err) } return info, nil } // WorkspaceInfo contains information about a workspace type WorkspaceInfo struct { Path string `json:"path"` FileCount int64 `json:"file_count"` PythonFiles int64 `json:"python_files"` NotebookFiles int64 `json:"notebook_files"` TextFiles int64 `json:"text_files"` ConfigFiles int64 `json:"config_files"` Size int64 `json:"size"` TotalSize int64 `json:"total_size"` Modified time.Time `json:"modified"` } // Cleanup removes unused mounts func (wm *WorkspaceManager) Cleanup() error { var toRemove []string for mountID, mount := range wm.mounts { if len(mount.Services) == 0 { toRemove = append(toRemove, mountID) } } for _, mountID := range toRemove { if err := wm.RemoveMount(mountID); err != nil { wm.logger.Warn("failed to remove unused mount", "mount_id", mountID, "error", err) } } if len(toRemove) > 0 { wm.logger.Info("cleanup completed", "removed_mounts", len(toRemove)) } return nil }