Skip to content

Commit

Permalink
Add the branch to init and pull commands
Browse files Browse the repository at this point in the history
  • Loading branch information
@moshe-kabala
moshe-kabala committed Nov 16, 2024
1 parent c5162cf commit bf719a9
Show file tree
Hide file tree
Showing 12 changed files with 195 additions and 44 deletions.
2 changes: 1 addition & 1 deletion cmd/code/create.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ func NewCreateCmd() *cobra.Command {
secretId = selectedSecret.GetCid()
}

err = workspace.CreateCodeTemplate(codeIntegration.GetCid(), "", secretId, codeIntegrationName)
err = workspace.CreateCodeTemplate(codeIntegration.GetCid(), "", secretId, codeIntegration.GetDefaultBranch(), codeIntegrationName)
if err != nil {
return err
}
Expand Down
21 changes: 14 additions & 7 deletions cmd/code/init.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ func init() {
var codeIntegrationId string
var newCodeIntegrationName string
var secretId string
var branch string

var cmd = &cobra.Command{
Use: "init",
Expand All @@ -28,8 +29,16 @@ func init() {
RunE: func(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()
var codeIntegration *code.CodeIntegration = nil
var err error
codeIntegrations, err := code.GetCodeIntegrations(ctx)
if err != nil {
return err
}

if len(newCodeIntegrationName) > 0 {
newCodeIntegrationName, err = code.AskForCodeIntegrationNameIfExisted(newCodeIntegrationName, codeIntegrations)
if err != nil {
return err
}
codeIntegration, err = code.AddCodeIntegration(ctx, newCodeIntegrationName)
if err != nil {
return err
Expand All @@ -48,29 +57,27 @@ func init() {
}

} else if len(codeIntegrationId) > 0 {
codeIntegrations, err := code.GetCodeIntegrations(ctx)
if err != nil {
return err
}

codeIntegration, err = entity.GetEntityById(codeIntegrationId, codeIntegrations, code.CodeIntegrationEntityDesc)
if err != nil {
return err
}
if len(secretId) == 0 {
latestVersion, err := code.GetLatestVersion(ctx, codeIntegration.GetCid())
latestVersion, err := code.GetLatestVersion(ctx, codeIntegration.GetCid(), branch)
if err != nil && err != code.ErrEmptyCodeIntegrationVersion {
return err
}
secretId = latestVersion.Metadata.GetSecretManagerId()
}
}
return workspace.CreateCodeTemplate(codeIntegration.GetCid(), "", secretId, ".")
return workspace.CreateCodeTemplate(codeIntegration.GetCid(), "", secretId, branch, ".")
},
}

cmd.Flags().StringVar(&codeIntegrationId, "codeId", "", "Code integration id of existing dataset")
cmd.Flags().StringVar(&newCodeIntegrationName, "new", "", "Name for new database")
cmd.Flags().StringVar(&secretId, "secretId", "", "Secret manager id for new code integration")
cmd.Flags().StringVarP(&branch, "branch", "b", "", "Branch for code integration")
cmd.MarkFlagsMutuallyExclusive("new", "codeId")
RootCommand.AddCommand(cmd)
}
18 changes: 13 additions & 5 deletions cmd/code/pull.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,23 +13,31 @@ import (

func NewPullCmd() *cobra.Command {
var cmd = &cobra.Command{
Use: "pull [dataset-id]",
Use: "pull [dataset-id] [branch]",
Short: "Pull dataset into a new directory",
Long: `Pull dataset into a new directory`,
RunE: func(cmd *cobra.Command, args []string) error {
var selectedDataset *code.CodeIntegration
var err error
var selectedBranch string
ctx := cmd.Context()
codeIntegrations, err := code.GetCodeIntegrations(ctx)
if err != nil {
return err
}
if len(args) == 0 {
selectedDataset, err = entity.SelectEntity(codeIntegrations, code.CodeIntegrationEntityDesc)

if err != nil {
return err
}
branches := code.BranchesFromCodeIntegration(selectedDataset)
selectedBranch, err = code.SelectBranch(branches, selectedDataset.DefaultBranch)
} else {
datasetId := args[0]
selectedDataset, err = entity.GetEntityById(datasetId, codeIntegrations, code.CodeIntegrationEntityDesc)
if args[1] != "" {
selectedBranch = args[1]
}
}
if err != nil {
return err
Expand All @@ -42,22 +50,22 @@ func NewPullCmd() *cobra.Command {
if dirExistsErr == nil {
return fmt.Errorf("can't pull '%s' dataset, directory named '%s' already exists on current directory", datasetName, datasetName)
}
latestVersion, err := code.GetLatestVersion(ctx, selectedDataset.GetCid())
latestVersion, err := code.GetLatestVersion(ctx, selectedDataset.GetCid(), selectedBranch)
secretId := latestVersion.Metadata.GetSecretManagerId()
if err == nil {
files, err := code.CloneCodeIntegrationVersion(ctx, latestVersion, datasetName)
if err != nil {
return err
}

workspaceConfig := workspace.NewWorkspaceConfig(selectedDataset.GetCid(), "", latestVersion.GetCodeEntryFile(), secretId, files)
workspaceConfig := workspace.NewWorkspaceConfig(selectedDataset.GetCid(), "", latestVersion.GetCodeEntryFile(), secretId, latestVersion.Branch, files)
err = workspace.SetWorkspaceConfig(workspaceConfig, datasetName)
if err != nil {
return err
}
} else if err == code.ErrEmptyCodeIntegrationVersion {
log.Warn("The selected dataset is empty, Create default template")
err = workspace.CreateCodeTemplate(selectedDataset.GetCid(), "", secretId, datasetName)
err = workspace.CreateCodeTemplate(selectedDataset.GetCid(), "", secretId, datasetName, selectedBranch)
if err != nil {
return err
}
Expand Down
10 changes: 9 additions & 1 deletion cmd/code/push.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import (

func NewPushCmd() *cobra.Command {
var secretId string
var branch string
var noWait bool
var force bool

Expand Down Expand Up @@ -45,7 +46,13 @@ func NewPushCmd() *cobra.Command {
return err
}

_, currentVersion, err := code.PushCode(ctx, force, codeIntegration.Cid, tarGzFile, workspaceConfig.EntryFile, secretId)
branches := code.BranchesFromCodeIntegration(codeIntegration)
branch, err := code.SyncBranchFromFlagAndConfig(branch, workspaceConfig, branches, codeIntegration.DefaultBranch)
if err != nil {
return err
}

_, currentVersion, err := code.PushCode(ctx, force, codeIntegration.Cid, tarGzFile, workspaceConfig.EntryFile, secretId, branch)
if err != nil {
return err
}
Expand All @@ -72,6 +79,7 @@ func NewPushCmd() *cobra.Command {
}

cmd.Flags().StringVar(&secretId, "secretId", "", "Secret id")
cmd.Flags().StringVarP(&branch, "branch", "b", "", "Branch")
cmd.Flags().BoolVar(&noWait, "no-wait", false, "Do not wait for code parsing")
cmd.Flags().BoolVarP(&force, "force", "f", false, "Force push code integration")

Expand Down
8 changes: 5 additions & 3 deletions cmd/models/import.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,8 @@ func NewImportCmd() *cobra.Command {
var projectId string
var message string
var modelType string
var branchName string
var modelBranch string
var codeIntegrationBranch string
var transformInput bool
var codeIntegrationId string
var noWait bool
Expand Down Expand Up @@ -57,7 +58,7 @@ func NewImportCmd() *cobra.Command {
projectId = selected.GetCid()
}

err = model.ImportModel(ctx, modelPath, projectId, message, modelType, branchName, codeIntegrationId, transformInput, !noWait)
err = model.ImportModel(ctx, modelPath, projectId, message, modelType, modelBranch, codeIntegrationId, codeIntegrationBranch, transformInput, !noWait)
if err != nil {
return err
}
Expand All @@ -68,7 +69,8 @@ func NewImportCmd() *cobra.Command {
cmd.Flags().StringVar(&projectId, "projectId", "", "ProjectId is the id of the project the model will be imported to")
cmd.Flags().StringVarP(&message, "message", "m", "", "Version message")
cmd.Flags().StringVar(&modelType, "type", "", "Type is the type of the model file [JSON_TF2 / ONNX / PB_TF2 / H5_TF2]")
cmd.Flags().StringVar(&branchName, "branch", "", "Branch is the name of the branch [OPTIONAL]")
cmd.Flags().StringVar(&modelBranch, "model-branch", "", "Name of the model branch [OPTIONAL]")
cmd.Flags().StringVar(&codeIntegrationBranch, "code-branch", "", "Name of the code integration branch [OPTIONAL]")
cmd.Flags().StringVar(&codeIntegrationId, "codeId", "", "This is a code integration id (Will use the last valid dataset version)")
cmd.Flags().BoolVar(&transformInput, "transform-input", true, "Transform input in case of ONNX model")
cmd.Flags().BoolVar(&noWait, "no-wait", false, "Do not wait for push to complete")
Expand Down
12 changes: 9 additions & 3 deletions cmd/projects/init.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ func NewInitCmd() *cobra.Command {
var projectId string
var secretId string
var codeIntegrationId string
var codeIntegrationBranch string

var cmd = &cobra.Command{
Use: "init",
Expand All @@ -35,26 +36,29 @@ func NewInitCmd() *cobra.Command {

isCreatingEmptyTemplate := true
if !wasCreatedCodeIntegration {
latestVersion, err := code.GetLatestVersion(ctx, codeIntegration.GetCid())
latestVersion, err := code.GetLatestVersion(ctx, codeIntegration.GetCid(), codeIntegrationBranch)
if err != nil && code.ErrEmptyCodeIntegrationVersion != err {
return err
} else if err == nil {
if secretId == "" {
secretId = latestVersion.Metadata.GetSecretManagerId()
}
isCreatingEmptyTemplate = false
_, err = code.CloneCodeIntegrationVersion(ctx, latestVersion, ".")
files, err := code.CloneCodeIntegrationVersion(ctx, latestVersion, ".")
if err != nil {
return err
}

err := workspace.OverrideWorkspaceConfig(
err = workspace.OverrideWorkspaceConfig(
codeIntegration.GetCid(),
selectedProject.GetCid(),
latestVersion.GetCodeEntryFile(),
secretId,
latestVersion.Branch,
files,
".",
)

if err != nil {
return err
}
Expand All @@ -76,6 +80,7 @@ func NewInitCmd() *cobra.Command {
codeIntegration.GetCid(),
selectedProject.GetCid(),
secretId,
codeIntegrationBranch,
".",
)
if err != nil {
Expand All @@ -89,6 +94,7 @@ func NewInitCmd() *cobra.Command {

cmd.Flags().StringVar(&projectId, "projectId", "", "ProjectId is the id of the project")
cmd.Flags().StringVar(&projectId, "codeId", "", "CodeIntegrationId is the id of the code integration to bind")
cmd.Flags().StringVar(&codeIntegrationBranch, "branch", "", "Branch of the code integration to bind")

return cmd
}
Expand Down
16 changes: 12 additions & 4 deletions cmd/projects/push.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,8 @@ func NewPushCmd() *cobra.Command {
var secretId string
var message string
var modelType string
var branchName string
var modelBranch string
var codeBranch string
var transformInput bool
var force bool
var noWait bool
Expand Down Expand Up @@ -83,6 +84,12 @@ func NewPushCmd() *cobra.Command {
return err
}

codeIntegrationBranches := code.BranchesFromCodeIntegration(codeIntegration)
codeBranch, err = code.SyncBranchFromFlagAndConfig(codeBranch, workspaceConfig, codeIntegrationBranches, codeIntegration.GetDefaultBranch())
if err != nil {
return err
}

secretId, err := secret.SyncSecretIdFromFlagAndConfig(ctx, secretId, workspaceConfig)
if err != nil {
return err
Expand All @@ -94,7 +101,7 @@ func NewPushCmd() *cobra.Command {
}
defer close()

pushed, currentVersion, err := code.PushCode(ctx, force, codeIntegration.Cid, tarGzFile, workspaceConfig.EntryFile, secretId)
pushed, currentVersion, err := code.PushCode(ctx, force, codeIntegration.Cid, tarGzFile, workspaceConfig.EntryFile, secretId, codeBranch)
if err != nil {
return err
}
Expand All @@ -119,7 +126,7 @@ func NewPushCmd() *cobra.Command {
return fmt.Errorf("latest code parsing failed, add --force to push anyway")
}

err = model.ImportModel(ctx, modelPath, currentProject.GetCid(), message, modelType, branchName, codeIntegration.GetCid(), transformInput, !noWait)
err = model.ImportModel(ctx, modelPath, currentProject.GetCid(), message, modelType, modelBranch, codeIntegration.GetCid(), codeBranch, transformInput, !noWait)
if err != nil {
return err
}
Expand All @@ -129,7 +136,8 @@ func NewPushCmd() *cobra.Command {

cmd.Flags().StringVarP(&message, "message", "m", "", "Version message")
cmd.Flags().StringVar(&modelType, "type", "", "Type is the type of the model file [JSON_TF2 / ONNX / PB_TF2 / H5_TF2]")
cmd.Flags().StringVar(&branchName, "branch", "", "Branch is the name of the branch [OPTIONAL]")
cmd.Flags().StringVar(&modelBranch, "model-branch", "", "Name of the model branch [OPTIONAL]")
cmd.Flags().StringVar(&codeBranch, "code-branch", "", "Name of the code branch [OPTIONAL]")
cmd.Flags().StringVar(&secretId, "secretId", "", "Secret id")
cmd.Flags().BoolVarP(&force, "force", "f", false, "Force push code integration")
cmd.Flags().BoolVar(&transformInput, "transform-input", true, "Transform input in case of ONNX model")
Expand Down
24 changes: 20 additions & 4 deletions pkg/code/api.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -73,9 +73,15 @@ func DeleteCodeIntegration(ctx context.Context, codeIntegration *CodeIntegration
return nil
}

func GetLatestVersion(ctx context.Context, codeIntegrationId string) (*tensorleapapi.DatasetVersion, error) {
func GetLatestVersion(ctx context.Context, codeIntegrationId string, branch string) (*tensorleapapi.DatasetVersion, error) {
params := *tensorleapapi.NewGetLatestDatasetVersionParams(codeIntegrationId)
if len(branch) == 0 {
branch = "master"
}
params.SetBranch(branch)

version, _, err := ApiClient.GetLatestDatasetVersion(ctx).
GetLatestDatasetVersionParams(*tensorleapapi.NewGetLatestDatasetVersionParams(codeIntegrationId)).
GetLatestDatasetVersionParams(params).
Execute()
if err != nil {
return nil, fmt.Errorf("failed to get latest version for code integration id: %s", codeIntegrationId)
Expand All @@ -91,7 +97,7 @@ func GetCodeIntegration(ctx context.Context, id string) (*CodeIntegrationVersion
return &res.DatasetVersion, err
}

func AddCodeIntegrationVersion(ctx context.Context, tarGzFile io.Reader, fileSize int64, codeIntegrationId, entryFile, secretId string) (*CodeIntegrationVersion, error) {
func AddCodeIntegrationVersion(ctx context.Context, tarGzFile io.Reader, fileSize int64, codeIntegrationId, entryFile, secretId, branch string) (*CodeIntegrationVersion, error) {

getDatasetVersionUploadUrlParams := *tensorleapapi.NewGetDatasetVersionUploadUrlParams(
codeIntegrationId,
Expand All @@ -116,6 +122,11 @@ func AddCodeIntegrationVersion(ctx context.Context, tarGzFile io.Reader, fileSiz
entryFile,
)

if len(branch) == 0 {
branch = "master"
}
saveDatasetVersionParams.SetBranch(branch)

if len(secretId) > 0 {
saveDatasetVersionParams.SecretManagerId = &secretId
}
Expand All @@ -127,8 +138,13 @@ func AddCodeIntegrationVersion(ctx context.Context, tarGzFile io.Reader, fileSiz
if err != nil {
return nil, err
}
for _, latestVersionPerBranch := range res.Dataset.LatestVersions {
if latestVersionPerBranch.Branch == branch {
return latestVersionPerBranch.Latest, nil
}
}

return res.Dataset.LatestVersion, nil
return nil, fmt.Errorf("failed to create a new code integration version")
}

func WaitForCodeIntegrationStatus(ctx context.Context, codeIntegrationId string) (ok bool, codeIntegrationVersion *CodeIntegrationVersion, err error) {
Expand Down
Loading

0 comments on commit bf719a9

Please sign in to comment.