feat(core): add support for Launch 🚀 (wandb#6822)

Co-authored-by: Kyle Goyette <[email protected]>
Co-authored-by: Dmitry Duev <[email protected]>

.circleci/config.yml

@@ -1478,6 +1478,18 @@ workflows:
           tox_args: "tests/pytest_tests/unit_tests/test_launch"
           codecov_args: "-F unit"
 
+      - tox-base:
+          requires:
+            - "unit-tests"
+          matrix:
+            parameters:
+              python_version_major: [3]
+              python_version_minor: [10]
+          name: "unit-launch-linux-py<<matrix.python_version_major>><<matrix.python_version_minor>>"
+          toxenv: "core-py<<matrix.python_version_major>><<matrix.python_version_minor>>"
+          tox_args: "tests/pytest_tests/unit_tests/test_launch"
+          codecov_args: "-F unit"
+
       #
       # Unit tests with pytest on Windows and MacOS
       #
@@ -1721,6 +1733,7 @@ workflows:
               python_version_major: [3]
               python_version_minor: [10]
               shard:
+                - "launch"
                 - "artifacts"
                 - "system_metrics"
           name: "system-<<matrix.shard>>-linux-py<<matrix.python_version_major>><<matrix.python_version_minor>>"

core/internal/data_types/data_types.go

@@ -1,9 +1,5 @@
 package data_types
 
-import (
-	"strings"
-)
-
 type TypeName string
 
 const (
@@ -46,7 +42,7 @@ type isParams_ParamsType interface {
 	isParams_ParamsType()
 }
 
-func GenerateTypeRepresentation(data interface{}) TypeRepresentation {
+func ResolveTypes(data interface{}) TypeRepresentation {
 	return resolveTypes(data)
 }
 
@@ -57,7 +53,7 @@ func resolveTypes(data interface{}, invalid ...bool) TypeRepresentation {
 	case map[string]interface{}:
 		result := make(map[string]TypeRepresentation)
 		for key, value := range v {
-			if !strings.HasPrefix(key, "_") {
+			if key != "_wandb" {
 				result[key] = resolveTypes(value)
 			}
 		}

core/internal/data_types/data_types_test.go

@@ -27,10 +27,9 @@ func TestGenerateTypeRepresentation(t *testing.T) {
 		{
 			name: "Nested Number, String, and Boolean",
 			input: map[string]interface{}{
-				"a":  1,
-				"b":  "hello",
-				"c":  true,
-				"_o": 42,
+				"a": 1,
+				"b": "hello",
+				"c": true,
 			},
 			expected: data_types.TypeRepresentation{
 				Name: data_types.MapTypeName,
@@ -186,10 +185,10 @@ func TestGenerateTypeRepresentation(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			result := data_types.GenerateTypeRepresentation(tc.input)
+			result := data_types.ResolveTypes(tc.input)
 			if !reflect.DeepEqual(result, tc.expected) {
 				jsonResult, _ := json.MarshalIndent(result, "", "  ")
-				t.Errorf("\nExpected: %v\nActual:   %v", tc.expected, jsonResult)
+				t.Errorf("\nExpected: %v\nActual:   %v", tc.expected, string(jsonResult))
 			}
 		})
 	}

core/pkg/artifacts/builder.go

@@ -3,6 +3,8 @@ package artifacts
 import (
 	"crypto/md5"
 	"fmt"
+	"io"
+	"os"
 	"sort"
 
 	"github.com/wandb/wandb/core/pkg/service"
@@ -40,7 +42,7 @@ func (b *ArtifactBuilder) initDefaultManifest() {
 }
 
 func (b *ArtifactBuilder) AddData(name string, dataMap map[string]interface{}) error {
-	filename, digest, err := utils.WriteJsonToFileWithDigest(dataMap)
+	filename, digest, size, err := utils.WriteJsonToFileWithDigest(dataMap)
 	if err != nil {
 		return err
 	}
@@ -49,6 +51,41 @@ func (b *ArtifactBuilder) AddData(name string, dataMap map[string]interface{}) e
 			Path:      name,
 			Digest:    digest,
 			LocalPath: filename,
+			Size:      size,
+		})
+	b.isDigestUpToDate = false
+	return nil
+}
+
+func (b *ArtifactBuilder) AddFile(path string, name string) error {
+
+	file, err := os.Open(path)
+	if err != nil {
+		return err
+	}
+
+	// file size:
+	var size int64
+	if stat, err := file.Stat(); err == nil { // if NO error
+		size = stat.Size()
+	}
+
+	defer file.Close()
+
+	data, err := io.ReadAll(file)
+	if err != nil {
+		return err
+	}
+	digest, err := utils.ComputeB64MD5(data)
+	if err != nil {
+		return err
+	}
+	b.artifactRecord.Manifest.Contents = append(b.artifactRecord.Manifest.Contents,
+		&service.ArtifactManifestEntry{
+			Path:      name,
+			Digest:    digest,
+			LocalPath: path,
+			Size:      size,
 		})
 	b.isDigestUpToDate = false
 	return nil

core/pkg/artifacts/builder_test.go

@@ -3,6 +3,7 @@ package artifacts
 import (
 	"encoding/json"
 	"fmt"
+	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -53,12 +54,24 @@ func TestArtifactBuilder(t *testing.T) {
 		ClientId:         clientId,
 		SequenceClientId: sequenceClientId,
 	}
+
 	builder := NewArtifactBuilder(baseArtifact)
 	err = builder.AddData("obj.object.json", weaveObjectData)
 	assert.Nil(t, err)
 	err = builder.AddData("obj.type.json", weaveTypeData)
 	assert.Nil(t, err)
+
+	file, err := os.CreateTemp("", "test-name.txt")
+	assert.Nil(t, err)
+	_, err = file.WriteString("wandb")
+	assert.Nil(t, err)
+	err = file.Sync()
+	assert.Nil(t, err)
+	file.Close()
+	err = builder.AddFile(file.Name(), "test-name.txt")
+	assert.Nil(t, err)
+
 	art := builder.GetArtifact()
-	assert.Equal(t, art.Digest, "62103abbae3f3d159ba71d1ffe37f2b1")
+	assert.Equal(t, art.Digest, "2f122f2bff8133c0d5806d9bac1b958c")
 	fmt.Printf("ART %+v\n", art)
 }

core/pkg/artifacts/manifest.go

@@ -63,7 +63,7 @@ func NewManifestFromProto(proto *service.ArtifactManifest) (Manifest, error) {
 	return manifest, nil
 }
 
-func (m *Manifest) WriteToFile() (filename string, digest string, rerr error) {
+func (m *Manifest) WriteToFile() (filename string, digest string, size int64, rerr error) {
 	return utils.WriteJsonToFileWithDigest(m)
 }
 

core/pkg/artifacts/saver.go

@@ -46,7 +46,9 @@ func NewArtifactSaver(
 }
 
 func (as *ArtifactSaver) createArtifact() (
-	attrs gql.CreateArtifactCreateArtifactCreateArtifactPayloadArtifact, rerr error) {
+	attrs gql.CreateArtifactCreateArtifactCreateArtifactPayloadArtifact,
+	rerr error,
+) {
 	aliases := []gql.ArtifactAliasInput{}
 	for _, alias := range as.Artifact.Aliases {
 		aliases = append(aliases,
@@ -382,7 +384,8 @@ func (as *ArtifactSaver) Save(ch chan<- *service.Record) (artifactID string, rer
 	if err != nil {
 		return "", fmt.Errorf("ArtifactSaver.resolveClientIDReferences: %w", err)
 	}
-	manifestFile, manifestDigest, err := manifest.WriteToFile()
+	// TODO: check if size is needed
+	manifestFile, manifestDigest, _, err := manifest.WriteToFile()
 	if err != nil {
 		return "", fmt.Errorf("ArtifactSaver.writeManifest: %w", err)
 	}