803677be57
- Add API server with WebSocket support and REST endpoints - Implement authentication system with API keys and permissions - Add task queue system with Redis backend and error handling - Include storage layer with database migrations and schemas - Add comprehensive logging, metrics, and telemetry - Implement security middleware and network utilities - Add experiment management and container orchestration - Include configuration management with smart defaults
105 lines
2.6 KiB
Go
105 lines
2.6 KiB
Go
// Package utils provides shared utilities for the fetch_ml project.
|
|
package container
|
|
|
|
import (
|
|
"fmt"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"strings"
|
|
|
|
"github.com/jfraeys/fetch_ml/internal/config"
|
|
)
|
|
|
|
// PodmanConfig holds configuration for Podman container execution
|
|
type PodmanConfig struct {
|
|
Image string
|
|
Workspace string
|
|
Results string
|
|
ContainerWorkspace string
|
|
ContainerResults string
|
|
GPUAccess bool
|
|
Memory string
|
|
CPUs string
|
|
}
|
|
|
|
// BuildPodmanCommand builds a Podman command for executing ML experiments
|
|
func BuildPodmanCommand(cfg PodmanConfig, scriptPath, requirementsPath string, extraArgs []string) *exec.Cmd {
|
|
args := []string{
|
|
"run", "--rm",
|
|
"--security-opt", "no-new-privileges",
|
|
"--cap-drop", "ALL",
|
|
}
|
|
|
|
if cfg.Memory != "" {
|
|
args = append(args, "--memory", cfg.Memory)
|
|
} else {
|
|
args = append(args, "--memory", config.DefaultPodmanMemory)
|
|
}
|
|
|
|
if cfg.CPUs != "" {
|
|
args = append(args, "--cpus", cfg.CPUs)
|
|
} else {
|
|
args = append(args, "--cpus", config.DefaultPodmanCPUs)
|
|
}
|
|
|
|
args = append(args, "--userns", "keep-id")
|
|
|
|
// Mount workspace
|
|
workspaceMount := fmt.Sprintf("%s:%s:rw", cfg.Workspace, cfg.ContainerWorkspace)
|
|
args = append(args, "-v", workspaceMount)
|
|
|
|
// Mount results
|
|
resultsMount := fmt.Sprintf("%s:%s:rw", cfg.Results, cfg.ContainerResults)
|
|
args = append(args, "-v", resultsMount)
|
|
|
|
if cfg.GPUAccess {
|
|
args = append(args, "--device", "/dev/dri")
|
|
}
|
|
|
|
// Image and command
|
|
args = append(args, cfg.Image,
|
|
"--workspace", cfg.ContainerWorkspace,
|
|
"--requirements", requirementsPath,
|
|
"--script", scriptPath,
|
|
)
|
|
|
|
// Add extra arguments via --args flag
|
|
if len(extraArgs) > 0 {
|
|
args = append(args, "--args")
|
|
args = append(args, extraArgs...)
|
|
}
|
|
|
|
return exec.Command("podman", args...)
|
|
}
|
|
|
|
// SanitizePath ensures a path is safe to use (prevents path traversal)
|
|
func SanitizePath(path string) (string, error) {
|
|
// Clean the path to remove any .. or . components
|
|
cleaned := filepath.Clean(path)
|
|
|
|
// Check for path traversal attempts
|
|
if strings.Contains(cleaned, "..") {
|
|
return "", fmt.Errorf("path traversal detected: %s", path)
|
|
}
|
|
|
|
return cleaned, nil
|
|
}
|
|
|
|
// ValidateJobName validates a job name is safe
|
|
func ValidateJobName(jobName string) error {
|
|
if jobName == "" {
|
|
return fmt.Errorf("job name cannot be empty")
|
|
}
|
|
|
|
// Check for dangerous characters
|
|
if strings.ContainsAny(jobName, "/\\<>:\"|?*") {
|
|
return fmt.Errorf("job name contains invalid characters: %s", jobName)
|
|
}
|
|
|
|
// Check for path traversal
|
|
if strings.Contains(jobName, "..") {
|
|
return fmt.Errorf("job name contains path traversal: %s", jobName)
|
|
}
|
|
|
|
return nil
|
|
}
|