Enforce usage of r8.im image name

pkg/docker/fast_push_test.go

@@ -22,7 +22,7 @@ import (
 func TestFastPush(t *testing.T) {
 	// Setup mock http server
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if r.URL.Path == "//test/versions" {
+		if r.URL.Path == "/username/modelname/versions" {
 			w.WriteHeader(http.StatusCreated)
 		} else {
 			w.WriteHeader(http.StatusConflict)
@@ -60,14 +60,14 @@ func TestFastPush(t *testing.T) {
 	monobeamClient := monobeam.NewClient(client)
 
 	// Run fast push
-	err = FastPush(context.Background(), "test", dir, command, webClient, monobeamClient)
+	err = FastPush(context.Background(), "r8.im/username/modelname", dir, command, webClient, monobeamClient)
 	require.NoError(t, err)
 }
 
 func TestFastPushWithWeight(t *testing.T) {
 	// Setup mock http server
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if r.URL.Path == "//test/versions" {
+		if r.URL.Path == "/username/modelname/versions" {
 			w.WriteHeader(http.StatusCreated)
 		} else {
 			w.WriteHeader(http.StatusConflict)
@@ -118,6 +118,6 @@ func TestFastPushWithWeight(t *testing.T) {
 	monobeamClient := monobeam.NewClient(client)
 
 	// Run fast push
-	err = FastPush(context.Background(), "test", dir, command, webClient, monobeamClient)
+	err = FastPush(context.Background(), "r8.im/username/modelname", dir, command, webClient, monobeamClient)
 	require.NoError(t, err)
 }

pkg/docker/push_test.go

@@ -20,7 +20,7 @@ import (
 func TestPush(t *testing.T) {
 	// Setup mock http server
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if r.URL.Path == "//test/versions" {
+		if r.URL.Path == "/username/modelname/versions" {
 			w.WriteHeader(http.StatusCreated)
 		} else {
 			w.WriteHeader(http.StatusConflict)
@@ -54,14 +54,14 @@ func TestPush(t *testing.T) {
 	command := dockertest.NewMockCommand()
 
 	// Run fast push
-	err = Push("test", true, dir, command)
+	err = Push("r8.im/username/modelname", true, dir, command)
 	require.NoError(t, err)
 }
 
 func TestPushWithWeight(t *testing.T) {
 	// Setup mock http server
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if r.URL.Path == "//test/versions" {
+		if r.URL.Path == "/username/modelname/versions" {
 			w.WriteHeader(http.StatusCreated)
 		} else {
 			w.WriteHeader(http.StatusConflict)
@@ -108,6 +108,6 @@ func TestPushWithWeight(t *testing.T) {
 	command := dockertest.NewMockCommand()
 
 	// Run fast push
-	err = Push("test", true, dir, command)
+	err = Push("r8.im/username/modelname", true, dir, command)
 	require.NoError(t, err)
 }

pkg/web/client.go

@@ -65,22 +65,21 @@ func NewClient(dockerCommand command.Command, client *http.Client) *Client {
 }
 
 func (c *Client) PostNewVersion(ctx context.Context, image string, weights []File, files []File) error {
-	userInfo, err := c.dockerCommand.LoadUserInformation(global.ReplicateRegistryHost)
+	version, err := c.versionFromManifest(image, weights, files)
 	if err != nil {
 		return err
 	}
 
-	version, err := c.versionFromManifest(image, weights, files)
+	jsonData, err := json.Marshal(version)
 	if err != nil {
 		return err
 	}
 
-	jsonData, err := json.Marshal(version)
+	versionUrl, err := newVersionURL(image)
 	if err != nil {
 		return err
 	}
 
-	versionUrl := newVersionURL(userInfo.Username, image)
 	req, err := http.NewRequestWithContext(ctx, http.MethodPost, versionUrl.String(), bytes.NewReader(jsonData))
 	if err != nil {
 		return err
@@ -192,10 +191,17 @@ func (c *Client) versionFromManifest(image string, weights []File, files []File)
 	return &version, nil
 }
 
-func newVersionURL(username string, image string) url.URL {
+func newVersionURL(image string) (url.URL, error) {
+	imageComponents := strings.Split(image, "/")
 	newVersionUrl := webBaseURL()
-	newVersionUrl.Path = strings.Join([]string{"", username, image, "versions"}, "/")
-	return newVersionUrl
+	if len(imageComponents) != 3 {
+		return newVersionUrl, errors.New("The image URL must have 3 components in the format of " + global.ReplicateRegistryHost + "/your-username/your-model")
+	}
+	if imageComponents[0] != global.ReplicateRegistryHost {
+		return newVersionUrl, errors.New("The image name must have the " + global.ReplicateRegistryHost + " prefix when using --x-fast.")
+	}
+	newVersionUrl.Path = strings.Join([]string{"", imageComponents[1], imageComponents[2], "versions"}, "/")
+	return newVersionUrl, nil
 }
 
 func webBaseURL() url.URL {

pkg/web/client_test.go

@@ -42,7 +42,7 @@ func TestPostNewVersion(t *testing.T) {
 
 	client := NewClient(command, http.DefaultClient)
 	ctx := context.Background()
-	err = client.PostNewVersion(ctx, "test", []File{}, []File{})
+	err = client.PostNewVersion(ctx, "r8.im/user/test", []File{}, []File{})
 	require.NoError(t, err)
 }
 
@@ -60,7 +60,7 @@ func TestVersionFromManifest(t *testing.T) {
 	dockertest.MockOpenAPISchema = "{\"test\": true}"
 
 	client := NewClient(command, http.DefaultClient)
-	version, err := client.versionFromManifest("test", []File{}, []File{})
+	version, err := client.versionFromManifest("r8.im/user/test", []File{}, []File{})
 	require.NoError(t, err)
 
 	var openAPISchema map[string]any
@@ -74,3 +74,13 @@ func TestVersionFromManifest(t *testing.T) {
 	require.Equal(t, openAPISchema, version.OpenAPISchema)
 	require.Equal(t, cogConfig, version.CogConfig)
 }
+
+func TestVersionURLErrorWithoutR8IMPrefix(t *testing.T) {
+	_, err := newVersionURL("docker.com/thing/thing")
+	require.Error(t, err)
+}
+
+func TestVersionURLErrorWithout3Components(t *testing.T) {
+	_, err := newVersionURL("username/test")
+	require.Error(t, err)
+}