diff --git a/cmd/run.go b/cmd/run.go index c5551ff..585f2fe 100644 --- a/cmd/run.go +++ b/cmd/run.go @@ -79,6 +79,7 @@ var ( runArgs string dataset string modelName string + taskOutput string ) func parseParams() (*models.RunConfig, error) { @@ -116,6 +117,7 @@ func parseParams() (*models.RunConfig, error) { Dataset: d, ModelName: m, RunArgs: rs.RunArgs, + TaskOutput: taskOutput, }, nil } @@ -143,10 +145,13 @@ func init() { runCmd.Flags().StringVarP(&modelName, "model_name", "n", "", "选择模型") + runCmd.Flags().StringVarP(&taskOutput, "task_output", "o", "", + "任务输出目录") // 设置必需参数 runCmd.MarkFlagRequired("resource") runCmd.MarkFlagRequired("image") runCmd.MarkFlagRequired("command") runCmd.MarkFlagRequired("dataset") + runCmd.MarkFlagRequired("task_output") } diff --git a/go.mod b/go.mod index 2c3176d..86ca7cd 100644 --- a/go.mod +++ b/go.mod @@ -4,9 +4,11 @@ go 1.24.5 require ( github.com/go-git/go-git/v5 v5.16.2 + github.com/mholt/archiver/v3 v3.5.1 github.com/nacos-group/nacos-sdk-go/v2 v2.3.2 github.com/pkg/errors v0.9.1 github.com/spf13/cobra v1.9.1 + golang.org/x/sync v0.13.0 ) require ( @@ -34,6 +36,7 @@ require ( github.com/aliyun/alibabacloud-dkms-transfer-go-sdk v0.1.8 // indirect github.com/aliyun/aliyun-secretsmanager-client-go v1.1.5 // indirect github.com/aliyun/credentials-go v1.4.3 // indirect + github.com/andybalholm/brotli v1.1.1 // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/buger/jsonparser v1.1.1 // indirect github.com/cespare/xxhash/v2 v2.2.0 // indirect @@ -41,21 +44,27 @@ require ( github.com/cloudflare/circl v1.6.1 // indirect github.com/cyphar/filepath-securejoin v0.4.1 // indirect github.com/deckarep/golang-set v1.7.1 // indirect + github.com/dsnet/compress v0.0.2-0.20230904184137-39efe44ab707 // indirect github.com/emirpasic/gods v1.18.1 // indirect github.com/go-git/gcfg v1.5.1-0.20230307220236-3a3c6141e376 // indirect github.com/go-git/go-billy/v5 v5.6.2 // indirect github.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8 // indirect github.com/golang/mock v1.6.0 // indirect github.com/golang/protobuf v1.5.4 // indirect + github.com/golang/snappy v0.0.4 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect github.com/jmespath/go-jmespath v0.0.0-20180206201540-c2b33e8439af // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/kevinburke/ssh_config v1.2.0 // indirect + github.com/klauspost/compress v1.17.11 // indirect + github.com/klauspost/pgzip v1.2.6 // indirect github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/nwaples/rardecode v1.1.0 // indirect github.com/orcaman/concurrent-map v0.0.0-20210501183033-44dafcb38ecc // indirect + github.com/pierrec/lz4/v4 v4.1.21 // indirect github.com/pjbgf/sha1cd v0.3.2 // indirect github.com/prometheus/client_golang v1.12.2 // indirect github.com/prometheus/client_model v0.2.0 // indirect @@ -65,13 +74,14 @@ require ( github.com/skeema/knownhosts v1.3.1 // indirect github.com/spf13/pflag v1.0.6 // indirect github.com/tjfoc/gmsm v1.4.1 // indirect + github.com/ulikunitz/xz v0.5.12 // indirect github.com/xanzy/ssh-agent v0.3.3 // indirect + github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8 // indirect go.uber.org/atomic v1.7.0 // indirect go.uber.org/multierr v1.6.0 // indirect go.uber.org/zap v1.21.0 // indirect golang.org/x/crypto v0.37.0 // indirect golang.org/x/net v0.39.0 // indirect - golang.org/x/sync v0.13.0 // indirect golang.org/x/sys v0.32.0 // indirect golang.org/x/text v0.24.0 // indirect golang.org/x/time v0.1.0 // indirect diff --git a/go.sum b/go.sum index 39b96f2..e824216 100644 --- a/go.sum +++ b/go.sum @@ -107,6 +107,9 @@ github.com/aliyun/credentials-go v1.3.6/go.mod h1:1LxUuX7L5YrZUWzBrRyk0SwSdH4OmP github.com/aliyun/credentials-go v1.3.10/go.mod h1:Jm6d+xIgwJVLVWT561vy67ZRP4lPTQxMbEYRuT2Ti1U= github.com/aliyun/credentials-go v1.4.3 h1:N3iHyvHRMyOwY1+0qBLSf3hb5JFiOujVSVuEpgeGttY= github.com/aliyun/credentials-go v1.4.3/go.mod h1:Jm6d+xIgwJVLVWT561vy67ZRP4lPTQxMbEYRuT2Ti1U= +github.com/andybalholm/brotli v1.0.1/go.mod h1:loMXtMfwqflxFJPmdbJO0a3KNoPuLBgiu3qAvBg8x/Y= +github.com/andybalholm/brotli v1.1.1 h1:PR2pgnyFznKEugtsUo0xLdDop5SKXd5Qf5ysW+7XdTA= +github.com/andybalholm/brotli v1.1.1/go.mod h1:05ib4cKhjx3OQYUY22hTVd34Bc8upXjOLL2rKwwZBoA= github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8= github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be/go.mod h1:ySMOLuWl6zY27l47sB3qLNK6tF2fkHG55UZxx8oIVo4= github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio= @@ -141,6 +144,10 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/deckarep/golang-set v1.7.1 h1:SCQV0S6gTtp6itiFrTqI+pfmJ4LN85S1YzhDf9rTHJQ= github.com/deckarep/golang-set v1.7.1/go.mod h1:93vsz/8Wt4joVM7c2AVqh+YRMiUSc14yDtF28KmMOgQ= +github.com/dsnet/compress v0.0.2-0.20210315054119-f66993602bf5/go.mod h1:qssHWj60/X5sZFNxpG4HBPDHVqxNm4DfnCKgrbZOT+s= +github.com/dsnet/compress v0.0.2-0.20230904184137-39efe44ab707 h1:2tV76y6Q9BB+NEBasnqvs7e49aEBFI8ejC89PSnWH+4= +github.com/dsnet/compress v0.0.2-0.20230904184137-39efe44ab707/go.mod h1:qssHWj60/X5sZFNxpG4HBPDHVqxNm4DfnCKgrbZOT+s= +github.com/dsnet/golib v0.0.0-20171103203638-1ea166775780/go.mod h1:Lj+Z9rebOhdfkVLjJ8T6VcRQv3SXugXy999NBtR9aFY= github.com/elazarl/goproxy v1.7.2 h1:Y2o6urb7Eule09PjlhQRGNsqRfPmYI3KKQLFpCAV3+o= github.com/elazarl/goproxy v1.7.2/go.mod h1:82vkLNir0ALaW14Rc399OTTjyNREgmdL2cVoIbS6XaE= github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc= @@ -204,6 +211,9 @@ github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaS github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/golang/snappy v0.0.2/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= +github.com/golang/snappy v0.0.4 h1:yAGX7huGHXlcLOEtBnF4w7FQwA26wojNCwOYAEhLjQM= +github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= github.com/google/btree v0.0.0-20180813153112-4030bb1f1f0c/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ= github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ= github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M= @@ -256,6 +266,14 @@ github.com/julienschmidt/httprouter v1.3.0/go.mod h1:JR6WtHb+2LUe8TCKY3cZOxFyyO8 github.com/kevinburke/ssh_config v1.2.0 h1:x584FjTGwHzMwvHx18PXxbBVzfnxogHaAReU4gf13a4= github.com/kevinburke/ssh_config v1.2.0/go.mod h1:CT57kijsi8u/K/BOFA39wgDQJ9CxiF4nAY/ojJ6r6mM= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= +github.com/klauspost/compress v1.4.1/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A= +github.com/klauspost/compress v1.11.4/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs= +github.com/klauspost/compress v1.17.11 h1:In6xLpyWOi1+C7tXUUWv2ot1QvBjxevKAaI6IXrJmUc= +github.com/klauspost/compress v1.17.11/go.mod h1:pMDklpSncoRMuLFrf1W9Ss9KT+0rH90U12bZKk7uwG0= +github.com/klauspost/cpuid v1.2.0/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek= +github.com/klauspost/pgzip v1.2.5/go.mod h1:Ch1tH69qFZu15pkjo5kYi6mth2Zzwzt50oCQKQE9RUs= +github.com/klauspost/pgzip v1.2.6 h1:8RXeL5crjEUFnR2/Sn6GJNWtSQ3Dk8pq4CL3jvdDyjU= +github.com/klauspost/pgzip v1.2.6/go.mod h1:Ch1tH69qFZu15pkjo5kYi6mth2Zzwzt50oCQKQE9RUs= github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= github.com/konsorten/go-windows-terminal-sequences v1.0.3/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc= @@ -268,6 +286,8 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/matttproud/golang_protobuf_extensions v1.0.1 h1:4hp9jkHxhMHkqkrB3Ix0jegS5sx/RkqARlsWZ6pIwiU= github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0= +github.com/mholt/archiver/v3 v3.5.1 h1:rDjOBX9JSF5BvoJGvjqK479aL70qh9DIpZCl+k7Clwo= +github.com/mholt/archiver/v3 v3.5.1/go.mod h1:e3dqJ7H78uzsRSEACH1joayhuSyhnonssnDhppzS1L4= github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= @@ -280,10 +300,15 @@ github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRW github.com/nacos-group/nacos-sdk-go/v2 v2.3.2 h1:9QB2nCJzT5wkTVlxNYl3XL/7+G6p2USMi2gQh/ouQQo= github.com/nacos-group/nacos-sdk-go/v2 v2.3.2/go.mod h1:9FKXl6FqOiVmm72i8kADtbeK71egyG9y3uRDBg41tpQ= github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno= +github.com/nwaples/rardecode v1.1.0 h1:vSxaY8vQhOcVr4mm5e8XllHWTiM4JF507A0Katqw7MQ= +github.com/nwaples/rardecode v1.1.0/go.mod h1:5DzqNKiOdpKKBH87u8VlvAnPZMXcGRhxWkRpHbbfGS0= github.com/onsi/gomega v1.34.1 h1:EUMJIKUjM8sKjYbtxQI9A4z2o+rruxnzNvpknOXie6k= github.com/onsi/gomega v1.34.1/go.mod h1:kU1QgUvBDLXBJq618Xvm2LUX6rSAfRaFRTcdOeDLwwY= github.com/orcaman/concurrent-map v0.0.0-20210501183033-44dafcb38ecc h1:Ak86L+yDSOzKFa7WM5bf5itSOo1e3Xh8bm5YCMUXIjQ= github.com/orcaman/concurrent-map v0.0.0-20210501183033-44dafcb38ecc/go.mod h1:Lu3tH6HLW3feq74c2GC+jIMS/K2CFcDWnWD9XkenwhI= +github.com/pierrec/lz4/v4 v4.1.2/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4= +github.com/pierrec/lz4/v4 v4.1.21 h1:yOVMLb6qSIDP67pl/5F7RepeKYu/VmTyEXvuMI5d9mQ= +github.com/pierrec/lz4/v4 v4.1.21/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4= github.com/pjbgf/sha1cd v0.3.2 h1:a9wb0bp1oC2TGwStyn0Umc/IGKQnEgF0vVaZ8QF8eo4= github.com/pjbgf/sha1cd v0.3.2/go.mod h1:zQWigSxVmsHEZow5qaLtPYxpcKMMQpa09ixqBxuCS6A= github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= @@ -347,8 +372,16 @@ github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf github.com/tjfoc/gmsm v1.3.2/go.mod h1:HaUcFuY0auTiaHB9MHFGCPx5IaLhTUd2atbCFBQXn9w= github.com/tjfoc/gmsm v1.4.1 h1:aMe1GlZb+0bLjn+cKTPEvvn9oUEBlJitaZiiBwsbgho= github.com/tjfoc/gmsm v1.4.1/go.mod h1:j4INPkHWMrhJb38G+J6W4Tw0AbuN8Thu3PbdVYhVcTE= +github.com/ulikunitz/xz v0.5.8/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14= +github.com/ulikunitz/xz v0.5.9/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14= +github.com/ulikunitz/xz v0.5.12 h1:37Nm15o69RwBkXM0J6A5OlE67RZTfzUxTj8fB3dfcsc= +github.com/ulikunitz/xz v0.5.12/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14= github.com/xanzy/ssh-agent v0.3.3 h1:+/15pJfg/RsTxqYcX6fHqOXZwwMP+2VyYWJeWM2qQFM= github.com/xanzy/ssh-agent v0.3.3/go.mod h1:6dzNDKs0J9rVPHPhaGCukekBHKqfl+L3KghI1Bc68Uw= +github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8 h1:nIPpBwaJSVYIxUFsDv3M8ofmx9yWTog9BfvIu0q41lo= +github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8/go.mod h1:HUYIGzjTL3rfEspMxjDjgmT5uz5wzYJKVo23qUhYTos= +github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= +github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.1.30/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= diff --git a/pkg/client/http_client.go b/pkg/client/http_client.go index 31461de..36c5481 100644 --- a/pkg/client/http_client.go +++ b/pkg/client/http_client.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "io" + "mime" "mime/multipart" "net/http" "net/url" @@ -212,6 +213,87 @@ func (c *HTTPClient) UploadFiles( return io.ReadAll(resp.Body) } +// DownloadFile 下载文件到本地路径（支持大文件流式处理） +// localPath：保存文件的本地路径（包含文件名） +// queryParams：可选查询参数 +func (c *HTTPClient) DownloadFile(path, localPath string, queryParams ...map[string]string) (string, error) { + // 构建完整URL + fullURL := c.BaseURL + path + + // 处理查询参数 + if len(queryParams) > 0 { + u, err := url.Parse(fullURL) + if err != nil { + return "", fmt.Errorf("解析URL失败: %w", err) + } + + q := u.Query() + for key, value := range queryParams[0] { + q.Add(key, value) + } + u.RawQuery = q.Encode() + fullURL = u.String() + } + + // 创建请求 + req, err := http.NewRequest("GET", fullURL, nil) + if err != nil { + return "", fmt.Errorf("创建请求失败: %w", err) + } + + // 设置请求头 + for k, v := range c.Headers { + req.Header.Set(k, v) + } + + // 发送请求 + resp, err := c.Client.Do(req) + if err != nil { + return "", fmt.Errorf("请求失败: %w", err) + } + defer resp.Body.Close() + + // 检查状态码 (接受2xx状态码) + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + // 尝试读取部分错误信息（限制长度防止内存溢出） + return "", &HTTPError{ + StatusCode: resp.StatusCode, + Message: resp.Status, + } + } + + filename := extractFilenameFromDisposition(resp.Header.Get("Content-Disposition")) + if filename == "" { + filename = filepath.Join(localPath, "file.zip") + } + // 创建本地文件 + outFile, err := os.Create(filename) + if err != nil { + return "", fmt.Errorf("创建本地文件失败: %w", err) + } + defer outFile.Close() + + // 流式拷贝（避免大文件内存溢出） + _, err = io.Copy(outFile, resp.Body) + if err != nil { + // 删除可能不完整的文件 + _ = os.Remove(localPath) + return "", fmt.Errorf("下载文件写入失败: %w", err) + } + + return filename, nil +} + +func extractFilenameFromDisposition(contentDisposition string) string { + // 尝试使用标准库的 mime.ParseMediaType + if _, params, err := mime.ParseMediaType(contentDisposition); err == nil { + if filename, exists := params["filename"]; exists && filename != "" { + return filename + } + } + return "" +} + // HTTPError 自定义HTTP错误 type HTTPError struct { StatusCode int diff --git a/pkg/models/params.go b/pkg/models/params.go index 8e5f97c..0987eed 100644 --- a/pkg/models/params.go +++ b/pkg/models/params.go @@ -14,6 +14,7 @@ type RunConfig struct { RunArgs []string `json:"run_args"` Dataset DataResourceConfig `json:"dataset"` ModelName DataResourceConfig `json:"model_name"` + TaskOutput string `json:"task_output"` } // ParseRunArgs 解析运行参数字符串为键值对 diff --git a/pkg/models/preparation.go b/pkg/models/preparation.go index 9686a95..3aaedce 100644 --- a/pkg/models/preparation.go +++ b/pkg/models/preparation.go @@ -202,3 +202,9 @@ type BindingResponse struct { type BindingResult struct { BindingID int `json:"bindingID"` } + +type BindResultSet struct { + BindCodeID int `json:"bind_code_id"` + BindDatasetID int `json:"bind_dataset_id"` + BindModelID int `json:"bind_model_id"` +} diff --git a/pkg/models/task.go b/pkg/models/task.go new file mode 100644 index 0000000..618fde3 --- /dev/null +++ b/pkg/models/task.go @@ -0,0 +1,256 @@ +package models + +import "time" + +// SubtaskRequest 表示子任务请求 +type SubtaskRequest struct { + UserID int `json:"userID"` + JobSetInfo JobSetInfo `json:"jobSetInfo"` +} + +// JobSetInfo 表示任务集信息 +type JobSetInfo struct { + Jobs []Job `json:"jobs"` +} + +// Job 表示单个任务 +type Job struct { + LocalJobID string `json:"localJobID"` + Name string `json:"name,omitempty"` + Description string `json:"description,omitempty"` + Type string `json:"type"` + Files *JobFiles `json:"files,omitempty"` + JobResources *JobResources `json:"jobResources,omitempty"` + TargetJob []TargetJob `json:"targetJob,omitempty"` +} + +// JobFiles 表示任务关联的文件 +type JobFiles struct { + Dataset FileBinding `json:"dataset"` + Model FileBinding `json:"model"` + Image ImageBinding `json:"image"` +} + +// FileBinding 表示文件绑定信息 +type FileBinding struct { + Type string `json:"type"` + BindingID int `json:"bindingID"` +} + +// ImageBinding 表示镜像绑定信息 +type ImageBinding struct { + Type string `json:"type"` + ImageID int `json:"imageID"` +} + +// JobResources 表示任务资源 +type JobResources struct { + ScheduleStrategy string `json:"scheduleStrategy"` + Clusters []Cluster `json:"clusters"` +} + +// Cluster 表示集群信息 +type Cluster struct { + ClusterID string `json:"clusterID"` + Runtime Runtime `json:"runtime"` + Code CodeInfo `json:"code"` + Resources []ResourcesItem `json:"resources"` +} + +// Runtime 表示运行时配置 +type Runtime struct { + Envs map[string]string `json:"envs"` + Params map[string]string `json:"params"` +} + +// CodeInfo 表示代码信息 +type CodeInfo struct { + Type string `json:"type"` + BindingID int `json:"bindingID"` +} + +// Resource 表示资源需求 +type Resource struct { + Type string `json:"type"` + Name string `json:"name"` + Number int `json:"number"` +} + +// TargetJob 表示目标任务 +type TargetJob struct { + TargetJobID string `json:"targetJobID"` + InputParams InputParams `json:"inputParams"` +} + +// InputParams 表示输入参数 +type InputParams struct { + PackageName string `json:"PackageName"` + ClusterID string `json:"ClusterID"` + Output string `json:"Output"` +} + +// SubmitTaskResponse 任务提交响应 +type SubmitTaskResponse struct { + Code string `json:"code"` // 状态码（字符串类型） + Message string `json:"message"` // 消息 + Data Data `json:"data"` // 响应数据 +} + +// Data 响应数据 +type Data struct { + JobSetID string `json:"jobSetID"` // 任务集ID + Message string `json:"message"` // 数据层消息 +} + +type TaskQueryRequest struct { + JobSetID string `json:"jobSetID"` // 用户ID + LocalJobID string `json:"localJobID"` // 页码 +} + +// TaskDetailResponse 任务详情查询响应 +type TaskDetailResponse struct { + Code string `json:"code"` // 状态码 + Message string `json:"message"` // 消息 + Data TaskDetailData `json:"data"` // 任务详情数据 +} + +// TaskDetailData 任务详情数据 +type TaskDetailData struct { + Name string `json:"name"` // 任务名称 + Description string `json:"description"` // 任务描述 + StartTime time.Time `json:"startTime"` // 开始时间 + EndTime time.Time `json:"endTime"` // 结束时间 + Strategy int `json:"strategy"` // 策略 + SynergyStatus int `json:"synergyStatus"` // 协同状态 + ClusterInfos []ClusterInfo `json:"clusterInfos"` // 集群信息列表 + SubTaskInfos []SubTaskInfo `json:"subTaskInfos"` // 子任务信息列表 + TaskTypeDict string `json:"taskTypeDict"` // 任务类型字典 + AdapterTypeDict string `json:"adapterTypeDict"` // 适配器类型字典 +} + +// ClusterInfo 集群信息 +type ClusterInfo struct { + ID string `json:"id"` // 集群ID + AdapterID string `json:"adapterId"` // 适配器ID + Name string `json:"name"` // 集群名称 + Nickname string `json:"nickname"` // 集群昵称 + Description string `json:"description"` // 集群描述 + CreateTime time.Time `json:"createTime"` // 创建时间 +} + +// SubTaskInfo 子任务信息 +type SubTaskInfo struct { + ID string `json:"id"` // 子任务ID + Name string `json:"name"` // 子任务名称 + ClusterID string `json:"clusterId"` // 集群ID + ClusterName string `json:"clusterName"` // 集群名称 + Status string `json:"status"` // 任务状态 + Remark string `json:"remark"` // 备注 + InferURL string `json:"inferUrl"` // 推理URL + WorkDir string `json:"workDir"` // 工作目录 + AppName string `json:"appName"` // 应用名称 +} + +// TaskLogResponse 任务日志响应 +type TaskLogResponse struct { + Code string `json:"code"` + Message string `json:"message"` + Data string `json:"data"` +} + +// TaskResultDetailResponse 任务详情响应 +type TaskResultDetailResponse struct { + Code string `json:"code"` // 状态码 + Message string `json:"message"` // 消息 + Data TaskResultDetailData `json:"data"` // 任务详情数据 +} + +// TaskResultDetailData 任务详情数据 +type TaskResultDetailData struct { + PcmJobData PcmJobData `json:"pcmJobData"` // PCM任务数据 +} + +// PcmJobData PCM任务数据 +type PcmJobData struct { + TaskID string `json:"taskID"` // 任务ID + JobSubmitInfo JobSubmitInfo `json:"jobSubmitInfo"` // 任务提交信息 + ResultFiles []ResultFile `json:"resultFiles"` // 结果文件列表 + Status string `json:"status"` // 任务状态 + ErrorMsg string `json:"errorMsg"` // 错误消息 +} + +// JobSubmitInfo 任务提交信息 +type JobSubmitInfo struct { + Type string `json:"type"` // 类型 + Info JobSubmitInfoDetail `json:"info"` // 任务详情 +} + +// JobSubmitInfoDetail 任务提交详情 +type JobSubmitInfoDetail struct { + Name string `json:"name"` // 任务名称 + Description string `json:"description"` // 任务描述 + JobResources JobResources `json:"jobResources"` // 任务资源 + DataDistributes DataDistributes `json:"dataDistributes"` // 数据分布 +} + +// DataDistributes 数据分布 +type DataDistributes struct { + Dataset []DatasetDistribute `json:"dataset"` // 数据集分布 + Code []CodeDistribute `json:"code"` // 代码分布 + Image []ImageDistribute `json:"image"` // 镜像分布 + Model []ModelDistribute `json:"model"` // 模型分布 +} + +// DatasetDistribute 数据集分布 +type DatasetDistribute struct { + DataName string `json:"dataName"` // 数据名称 + PackageID int `json:"packageID"` // 包ID + Clusters []ClusterStorage `json:"clusters"` // 集群存储 +} + +// CodeDistribute 代码分布 +type CodeDistribute struct { + DataName string `json:"dataName"` // 数据名称 + PackageID int `json:"packageID"` // 包ID + Output string `json:"output"` // 输出路径 + Clusters []ClusterStorage `json:"clusters"` // 集群存储 +} + +// ImageDistribute 镜像分布 +type ImageDistribute struct { + DataName string `json:"dataName"` // 数据名称 + PackageID int `json:"packageID"` // 包ID + Clusters []ClusterStorage `json:"clusters"` // 集群存储 +} + +// ModelDistribute 模型分布 +type ModelDistribute struct { + DataName string `json:"dataName"` // 数据名称 + PackageID int `json:"packageID"` // 包ID + Clusters []ClusterStorage `json:"clusters"` // 集群存储 +} + +// ClusterStorage 集群存储信息 +type ClusterStorage struct { + ClusterID string `json:"clusterID"` // 集群ID + StorageID int `json:"storageID"` // 存储ID + JsonData string `json:"jsonData"` // JSON数据 +} + +// ResultFile 结果文件 +type ResultFile struct { + ClusterID string `json:"clusterID"` // 集群ID + Objects []ResultObject `json:"objects"` // 结果对象列表 +} + +// ResultObject 结果对象 +type ResultObject struct { + ObjectID int `json:"objectID"` // 对象ID + PackageID int `json:"packageID"` // 包ID + Path string `json:"path"` // 路径 + Size string `json:"size"` // 大小 + FileHash string `json:"fileHash"` // 文件哈希 + Redundancy Redundancy `json:"redundancy"` // 冗余配置 + CreateTime time.Time `json:"createTime"` // 创建时间 + UpdateTime time.Time `json:"updateTime"` // 更新时间 +} diff --git a/pkg/service/interface.go b/pkg/service/interface.go index 876d53b..96c91e7 100644 --- a/pkg/service/interface.go +++ b/pkg/service/interface.go @@ -11,11 +11,17 @@ type AuthService interface { } type PreparationService interface { + PrepareAll(ctx context.Context, authData models.AuthData, config models.RunConfig, + clusterID string) (*models.BindResultSet, error) PrepareCode(ctx context.Context, authData models.AuthData, config models.RunConfig, clientID string) (int, error) PrepareDataset(ctx context.Context, authData models.AuthData, config models.DataResourceConfig, clientID string) (int, error) PrepareModel(ctx context.Context, authData models.AuthData, config models.DataResourceConfig, clientID string) (int, error) } type TaskService interface { - CreateTask(ctx context.Context) (int, error) + SubmitTask(ctx context.Context, authData models.AuthData, config models.RunConfig, + clusterID string, bindResultSet models.BindResultSet) (string, error) + GetTaskStatus(ctx context.Context, authData models.AuthData, jobSetID string) (string, error) + GetTaskLogs(ctx context.Context, authData models.AuthData, jobSetID string) (string, error) + GetTaskResult(ctx context.Context, authData models.AuthData, path string, jobSetID string) error } diff --git a/pkg/service/preparation/preparation_service.go b/pkg/service/preparation/preparation_service.go index 15f009d..4ddce26 100644 --- a/pkg/service/preparation/preparation_service.go +++ b/pkg/service/preparation/preparation_service.go @@ -9,6 +9,7 @@ import ( "fmt" "github.com/go-git/go-git/v5" "github.com/go-git/go-git/v5/plumbing" + "golang.org/x/sync/errgroup" "net/url" "os" "path/filepath" @@ -24,6 +25,7 @@ const ( SceneType = "sceneAPI" DatasetType = "dataset" ModelType = "model" + ImageType = "image" CodeType = "code" BindingType = "binding" ImageCategory = "image" @@ -42,6 +44,84 @@ func NewPreparationService(baseURL string, timeout time.Duration) service.Prepar } } +func (s *preparationService) PrepareAll(ctx context.Context, authData models.AuthData, config models.RunConfig, + clusterID string) (*models.BindResultSet, error) { + // 创建错误组和上下文 + errGroup, ctx := errgroup.WithContext(ctx) + + // 结果通道 + type result struct { + id int + err error + } + + codeChan := make(chan result, 1) + datasetChan := make(chan result, 1) + modelChan := make(chan result, 1) + + // 并发执行代码准备 + errGroup.Go(func() error { + id, err := s.PrepareCode(ctx, authData, config, clusterID) + codeChan <- result{id, err} + return err + }) + + // 并发执行数据集准备 + errGroup.Go(func() error { + id, err := s.PrepareDataset(ctx, authData, config.Dataset, clusterID) + datasetChan <- result{id, err} + return err + }) + + // 并发执行模型准备 + errGroup.Go(func() error { + id, err := s.PrepareModel(ctx, authData, config.ModelName, clusterID) + modelChan <- result{id, err} + return err + }) + + // 等待所有任务完成 + if err := errGroup.Wait(); err != nil { + // 有错误发生，但我们需要收集所有结果 + codeRes := <-codeChan + datasetRes := <-datasetChan + modelRes := <-modelChan + + return &models.BindResultSet{ + BindCodeID: codeRes.id, + BindDatasetID: datasetRes.id, + BindModelID: modelRes.id, + }, err + } + + // 所有任务成功完成 + codeRes := <-codeChan + datasetRes := <-datasetChan + modelRes := <-modelChan + + return &models.BindResultSet{ + BindCodeID: codeRes.id, + BindDatasetID: datasetRes.id, + BindModelID: modelRes.id, + }, nil +} + +// 聚合多个错误 +func aggregateErrors(errs ...error) error { + var errMsgs []string + for _, err := range errs { + if err != nil { + errMsgs = append(errMsgs, err.Error()) + } + } + + if len(errMsgs) == 0 { + return nil + } + + return fmt.Errorf("多个错误发生: %s", strings.Join(errMsgs, "; ")) +} + func (s *preparationService) PrepareCode(ctx context.Context, authData models.AuthData, config models.RunConfig, clusterID string) (int, error) { resul, err := s.uploadCode(ctx, authData, config.CodeConfig) @@ -51,6 +131,26 @@ func (s *preparationService) PrepareCode(ctx context.Context, authData models.Au return s.bindCode(ctx, authData, config, resul, clusterID) } +func (s *preparationService) PrepareDataset(ctx context.Context, authData models.AuthData, config models.DataResourceConfig, + clusterID string) (int, error) { + // 1 上传数据集到数据集仓库 + result, err := s.uploadFiles(ctx, authData, config, DatasetType) + if err != nil { + return 0, fmt.Errorf("上传数据集失败: %w", err) + } + //2 绑定数据集到用户的bucket + return s.bindDataset(ctx, authData, config, result, clusterID) +} + +func (s *preparationService) PrepareModel(ctx context.Context, authData models.AuthData, config models.DataResourceConfig, + clusterID string) (int, error) { + result, err := s.uploadFiles(ctx, authData, config, ModelType) + if err != nil { + return 0, fmt.Errorf("上传模型失败: %w", err) + } + return s.bindModel(ctx, authData, config, result, clusterID) +} + func (s *preparationService) bindCode(ctx context.Context, authData models.AuthData, config models.RunConfig, uploadResult *models.PackageCreateResponse, clusterID string) (int, error) { appInstanceName := generateAppInstanceName(CodeType) @@ -103,26 +203,6 @@ func (s *preparationService) getStartCommandObjectID(config models.RunConfig, up return 0, fmt.Errorf("未找到启动命令文件: %s", config.Command) } -func (s *preparationService) PrepareDataset(ctx context.Context, authData models.AuthData, config models.DataResourceConfig, - clusterID string) (int, error) { - // 1 上传数据集到数据集仓库 - result, err := s.uploadFiles(ctx, authData, config, DatasetType) - if err != nil { - return 0, fmt.Errorf("上传数据集失败: %w", err) - } - //2 绑定数据集到用户的bucket - return s.bindDataset(ctx, authData, config, result, clusterID) -} - -func (s *preparationService) PrepareModel(ctx context.Context, authData models.AuthData, config models.DataResourceConfig, - clusterID string) (int, error) { - result, err := s.uploadFiles(ctx, authData, config, ModelType) - if err != nil { - return 0, fmt.Errorf("上传模型失败: %w", err) - } - return s.bindModel(ctx, authData, config, result, clusterID) -} - func (s *preparationService) bindModel(ctx context.Context, authData models.AuthData, config models.DataResourceConfig, uploadResult *models.PackageCreateResponse, clusterID string) (int, error) { appInstanceName := generateAppInstanceName(ModelType) diff --git a/pkg/service/task/task_service.go b/pkg/service/task/task_service.go new file mode 100644 index 0000000..26df10e --- /dev/null +++ b/pkg/service/task/task_service.go @@ -0,0 +1,239 @@ +package task + +import ( + "context" + "encoding/json" + "fmt" + "github.com/mholt/archiver/v3" + "os" + "remote-task-excutor-cli/pkg/client" + "remote-task-excutor-cli/pkg/models" + "remote-task-excutor-cli/pkg/service" + "remote-task-excutor-cli/pkg/service/preparation" + "time" +) + +const ( + MainTaskID = "1" + DataReturnTaskID = "4" + AITaskType = "AI" + DataReturnType = "DataReturn" + + JobSetIDKey = "jobSetID" + LocalJobIDKey = "localJobID" + ResponseOK = "OK" +) + +type taskService struct { + httpClient *client.HTTPClient +} + +func NewTaskServiceService(baseURL string, timeout time.Duration) service.TaskService { + return &taskService{ + httpClient: client.NewHTTPClient(baseURL, timeout), + } +} + +func (s *taskService) SubmitTask(ctx context.Context, authData models.AuthData, config models.RunConfig, + clusterID string, bindResultSet models.BindResultSet) (string, error) { + submitTaskReq := s.buildSubmitTaskReq(ctx, authData, config, clusterID, bindResultSet) + s.httpClient.Headers["Authorization"] = "Bearer " + authData.Token + resp, err := s.httpClient.PostJSON("/jsm/jobSet/submit", submitTaskReq) + if err != nil { + fmt.Println("Submit task failed: ", err) + return "", err + } + + var submitTaskResp models.SubmitTaskResponse + if err := json.Unmarshal(resp, &submitTaskResp); err != nil { + fmt.Println("Submit task response unmarshal failed: ", err) + return "", err + } + if submitTaskResp.Code != ResponseOK { + fmt.Println("Submit task failed: ", submitTaskResp.Code) + return "", fmt.Errorf("submit task failed: %s", submitTaskResp.Code) + } + fmt.Println("Submit task result: ", resp) + return submitTaskResp.Data.JobSetID, nil +} + +func (s *taskService) GetTaskStatus(ctx context.Context, authData models.AuthData, jobSetID string) (string, error) { + s.httpClient.Headers["Authorization"] = "Bearer " + authData.Token + params := map[string]string{ + JobSetIDKey: jobSetID, + LocalJobIDKey: MainTaskID, + } + resp, err := s.httpClient.Get("/jsm/jobMgr/detail", params) + if err != nil { + fmt.Println("Get task status failed: ", err) + return "", err + } + + var taskStatus models.TaskDetailResponse + if err := json.Unmarshal(resp, &taskStatus); err != nil { + fmt.Println("Get task status response unmarshal failed: ", err) + return "", err + } + + if taskStatus.Code != ResponseOK { + fmt.Println("Get task status failed: ", taskStatus.Code) + return "", fmt.Errorf("get task status failed: code: %s, message: %s", taskStatus.Code, taskStatus.Message) + } + fmt.Println("Get task status result: ", resp) + return taskStatus.Data.SubTaskInfos[0].Status, nil +} + +func (s *taskService) GetTaskLogs(ctx context.Context, authData models.AuthData, jobSetID string) (string, error) { + s.httpClient.Headers["Authorization"] = "Bearer " + authData.Token + params := map[string]string{ + JobSetIDKey: jobSetID, + LocalJobIDKey: MainTaskID, + } + resp, err := s.httpClient.Get("/jsm/jobMgr/log", params) + if err != nil { + fmt.Println("Get task logs failed: ", err) + return "", err + } + + var taskLogResponse models.TaskLogResponse + if err := json.Unmarshal(resp, &taskLogResponse); err != nil { + fmt.Println("Get task logs response unmarshal failed: ", err) + return "", err + } + + if taskLogResponse.Code != ResponseOK { + fmt.Println("Get task logs failed: ", err) + return "", fmt.Errorf("get task logs failed: %s", taskLogResponse.Code) + } + + fmt.Println("Get task logs result: ", resp) + return taskLogResponse.Data, nil +} + +func (s *taskService) GetTaskResult(ctx context.Context, authData models.AuthData, path string, jobSetID string) error { + result, err := s.getTaskResultDetail(ctx, authData, jobSetID) + if err != nil { + return err + } + + filename, err := s.downloadTaskResult(ctx, authData, path, result.PcmJobData.ResultFiles[0].Objects[0].PackageID) + if err != nil { + return err + } + + if err := archiver.Unarchive(filename, path); err != nil { + fmt.Println("Unarchive task result failed: ", err) + return err + } + // 删除filename文件 + if err := os.Remove(filename); err != nil { + fmt.Println("Remove task result failed: ", err) + return err + } + + return nil +} + +func (s *taskService) downloadTaskResult(ctx context.Context, authData models.AuthData, path string, packageID int) (string, error) { + s.httpClient.Headers["Authorization"] = "Bearer " + authData.Token + params := map[string]string{ + "packageId": fmt.Sprintf("%s", packageID), + "userId": fmt.Sprintf("%s", authData.JsmUserInfo.Data.UserID), + } + filename, err := s.httpClient.DownloadFile("/jcs/v1/package/download", path, params) + if err != nil { + fmt.Println("Download task result failed: ", err) + return "", err + } + return filename, nil +} + +func (s *taskService) getTaskResultDetail(ctx context.Context, authData models.AuthData, + jobSetID string) (*models.TaskResultDetailData, error) { + s.httpClient.Headers["Authorization"] = "Bearer " + authData.Token + queryTaskResultReq := models.TaskQueryRequest{ + JobSetID: jobSetID, + LocalJobID: DataReturnTaskID, + } + + resp, err := s.httpClient.PostJSON("/jsm/jobMgr/result", queryTaskResultReq) + if err != nil { + fmt.Println("Get task result details failed: ", err) + return nil, err + } + + var taskResultDetails models.TaskResultDetailResponse + if err := json.Unmarshal(resp, &taskResultDetails); err != nil { + fmt.Println("Get task result details response unmarshal failed: ", err) + return nil, err + } + + if taskResultDetails.Code != ResponseOK { + fmt.Println("Get task result details failed: ", taskResultDetails.Code) + return nil, fmt.Errorf("get task result details failed: %s", taskResultDetails.Code) + } + + return &taskResultDetails.Data, nil +} + +func (s *taskService) buildSubmitTaskReq(ctx context.Context, authData models.AuthData, config models.RunConfig, + clusterID string, bindResultSet models.BindResultSet) models.SubtaskRequest { + return models.SubtaskRequest{ + UserID: authData.JsmUserInfo.Data.UserID, + JobSetInfo: models.JobSetInfo{ + Jobs: []models.Job{ + { + LocalJobID: MainTaskID, + Name: "remote-task" + preparation.GenerateUniqueID(), + Description: "材料平台远程任务", + Type: AITaskType, + Files: &models.JobFiles{ + Dataset: models.FileBinding{ + Type: preparation.DatasetType, + BindingID: bindResultSet.BindDatasetID, + }, + Model: models.FileBinding{ + Type: preparation.ModelType, + BindingID: bindResultSet.BindModelID, + }, + Image: models.ImageBinding{ + Type: preparation.ImageType, + ImageID: config.Image, + }, + }, + JobResources: &models.JobResources{ + ScheduleStrategy: "dataLocality", + Clusters: []models.Cluster{ + { + ClusterID: clusterID, + Runtime: models.Runtime{ + Envs: nil, + Params: nil, + }, + Code: models.CodeInfo{ + Type: preparation.CodeType, + BindingID: bindResultSet.BindCodeID, + }, + Resources: config.Resource.Resources, + }, + }, + }, + }, + { + LocalJobID: DataReturnTaskID, + Type: DataReturnType, + TargetJob: []models.TargetJob{ + { + TargetJobID: MainTaskID, + InputParams: models.InputParams{ + PackageName: "name", + ClusterID: "ClusterID", + Output: "Output", + }, + }, + }, + }, + }, + }, + } +} diff --git a/pkg/service/task/task_service_test.go b/pkg/service/task/task_service_test.go new file mode 100644 index 0000000..5ed7c61 --- /dev/null +++ b/pkg/service/task/task_service_test.go @@ -0,0 +1,29 @@ +package task + +import "testing" + +func Test_extractFilenameFromDisposition(t *testing.T) { + type args struct { + contentDisposition string + } + tests := []struct { + name string + args args + want string + }{ + { + name: "extract filename from content-disposition", + args: args{ + contentDisposition: "attachment; filename=trainingtask-kdsrdtueyyxi_20250801095233.zip", + }, + want: "trainingtask-kdsrdtueyyxi_20250801095233.zip", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := extractFilenameFromDisposition(tt.args.contentDisposition); got != tt.want { + t.Errorf("extractFilenameFromDisposition() = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/remote-task-excutor-cli.exe b/remote-task-excutor-cli.exe index d539c4e..911ca72 100644 Binary files a/remote-task-excutor-cli.exe and b/remote-task-excutor-cli.exe differ