Fix cortex deploy panic when object and prefix present on bucket

pkg/cortex/serve/cortex_internal/lib/storage/gcs.py

@@ -57,15 +57,30 @@ def is_valid_gcs_path(path: str) -> bool:
     def _does_blob_exist(self, prefix: str) -> bool:
         return isinstance(self.gcs.get_blob(blob_name=prefix), storage.blob.Blob)
 
+    def _gcs_matching_blobs_generator(self, max_results=None, prefix="", include_dir_objects=False):
+        for blob in self.gcs.list_blobs(max_results, prefix=prefix):
+            if include_dir_objects or not blob.name.endswith("/"):
+                yield blob
+
     def _is_gcs_dir(self, dir_path: str) -> bool:
         prefix = util.ensure_suffix(dir_path, "/")
-        return len(list(self.gcs.list_blobs(max_results=2, prefix=prefix))) > 1
+        matching_blobs = list(
+            self._gcs_matching_blobs_generator(
+                max_results=2, prefix=prefix, include_dir_objects=True
+            )
+        )
+
+        return len(matching_blobs) > 0
 
-    def search(self, prefix: str = "") -> Tuple[List[str], List[datetime.datetime]]:
+    def search(
+        self, prefix: str = "", include_dir_objects=False
+    ) -> Tuple[List[str], List[datetime.datetime]]:
         paths = []
         timestamps = []
 
-        for blob in self.gcs.list_blobs(prefix=prefix):
+        for blob in self._gcs_matching_blobs_generator(
+            prefix=prefix, include_dir_objects=include_dir_objects
+        ):
             paths.append(blob.name)
             timestamps.append(blob.updated)
 
@@ -112,12 +127,14 @@ def download_dir(self, prefix: str, local_dir: str):
     def download_dir_contents(self, prefix: str, local_dir: str):
         util.mkdir_p(local_dir)
         prefix = util.ensure_suffix(prefix, "/")
-        for blob in self.gcs.list_blobs(prefix=prefix):
-            if blob.name.endswith("/"):
-                continue
+        for blob in self._gcs_matching_blobs_generator(prefix=prefix, include_dir_objects=True):
             relative_path = util.trim_prefix(blob.name, prefix)
             local_dest_path = os.path.join(local_dir, relative_path)
-            self.download_file(blob.name, local_dest_path)
+
+            if not local_dest_path.endswith("/"):
+                self.download_file(blob.name, local_dest_path)
+            else:
+                util.mkdir_p(os.path.dirname(local_dest_path))
 
     def download(self, prefix: str, local_dir: str):
         """

pkg/cortex/serve/cortex_internal/lib/storage/s3.py

@@ -85,7 +85,7 @@ def _is_s3_dir(self, dir_path):
         prefix = util.ensure_suffix(dir_path, "/")
         return self._is_s3_prefix(prefix)
 
-    def _get_matching_s3_objects_generator(self, prefix="", suffix=""):
+    def _get_matching_s3_objects_generator(self, prefix="", suffix="", include_dir_objects=False):
         kwargs = {"Bucket": self.bucket, "Prefix": prefix}
 
         while True:
@@ -97,16 +97,20 @@ def _get_matching_s3_objects_generator(self, prefix="", suffix=""):
 
             for obj in contents:
                 key = obj["Key"]
-                if key.startswith(prefix) and key.endswith(suffix):
+                if (
+                    key.startswith(prefix)
+                    and key.endswith(suffix)
+                    and (include_dir_objects or not key.endswith("/"))
+                ):
                     yield obj
 
             try:
                 kwargs["ContinuationToken"] = resp["NextContinuationToken"]
             except KeyError:
                 break
 
-    def _get_matching_s3_keys_generator(self, prefix="", suffix=""):
-        for obj in self._get_matching_s3_objects_generator(prefix, suffix):
+    def _get_matching_s3_keys_generator(self, prefix="", suffix="", include_dir_objects=False):
+        for obj in self._get_matching_s3_objects_generator(prefix, suffix, include_dir_objects):
             yield obj["Key"], obj["LastModified"]
 
     def put_object(self, body, key):
@@ -146,18 +150,15 @@ def _read_bytes_from_s3_single(self, key, allow_missing=False, ext_bucket=None):
 
         return byte_array.strip()
 
-    def search(self, prefix="", suffix="") -> Tuple[List[str], List[datetime.datetime]]:
+    def search(
+        self, prefix="", suffix="", include_dir_objects=False
+    ) -> Tuple[List[str], List[datetime.datetime]]:
         paths = []
         timestamps = []
 
-        timestamp_map = {}
-        for key, ts in self._get_matching_s3_keys_generator(prefix, suffix):
-            timestamp_map[key] = ts
-
-        filtered_keys = util.remove_non_empty_directory_paths(list(timestamp_map.keys()))
-        for key in filtered_keys:
+        for key, ts in self._get_matching_s3_keys_generator(prefix, suffix, include_dir_objects):
             paths.append(key)
-            timestamps.append(timestamp_map[key])
+            timestamps.append(ts)
 
         return paths, timestamps
 
@@ -217,12 +218,14 @@ def download_dir(self, prefix, local_dir):
     def download_dir_contents(self, prefix, local_dir):
         util.mkdir_p(local_dir)
         prefix = util.ensure_suffix(prefix, "/")
-        for key, _ in self._get_matching_s3_keys_generator(prefix):
-            if key.endswith("/"):
-                continue
+        for key, _ in self._get_matching_s3_keys_generator(prefix, include_dir_objects=True):
             rel_path = util.trim_prefix(key, prefix)
             local_dest_path = os.path.join(local_dir, rel_path)
-            self.download_file(key, local_dest_path)
+
+            if not local_dest_path.endswith("/"):
+                self.download_file(key, local_dest_path)
+            else:
+                util.mkdir_p(os.path.dirname(local_dest_path))
 
     def download_and_unzip(self, key, local_dir):
         util.mkdir_p(local_dir)

pkg/lib/gcp/gcs.go

@@ -255,12 +255,12 @@ func ConvertGCSObjectsToKeys(gcsObjects ...*storage.ObjectAttrs) []string {
 	return paths
 }
 
-func (c *Client) ListGCSDir(bucket string, gcsDir string, maxResults *int64) ([]*storage.ObjectAttrs, error) {
+func (c *Client) ListGCSDir(bucket string, gcsDir string, includeDirObjects bool, maxResults *int64) ([]*storage.ObjectAttrs, error) {
 	gcsDir = s.EnsureSuffix(gcsDir, "/")
-	return c.ListGCSPrefix(bucket, gcsDir, maxResults)
+	return c.ListGCSPrefix(bucket, gcsDir, includeDirObjects, maxResults)
 }
 
-func (c *Client) ListGCSPrefix(bucket string, prefix string, maxResults *int64) ([]*storage.ObjectAttrs, error) {
+func (c *Client) ListGCSPrefix(bucket string, prefix string, includeDirObjects bool, maxResults *int64) ([]*storage.ObjectAttrs, error) {
 	gcsClient, err := c.GCS()
 	if err != nil {
 		return nil, err
@@ -282,7 +282,11 @@ func (c *Client) ListGCSPrefix(bucket string, prefix string, maxResults *int64)
 		if attrs == nil {
 			continue
 		}
-		gcsObjects = append(gcsObjects, attrs)
+
+		if includeDirObjects || !strings.HasSuffix(attrs.Name, "/") {
+			gcsObjects = append(gcsObjects, attrs)
+		}
+
 		if maxResults != nil && int64(len(gcsObjects)) >= *maxResults {
 			break
 		}
@@ -291,12 +295,12 @@ func (c *Client) ListGCSPrefix(bucket string, prefix string, maxResults *int64)
 	return gcsObjects, nil
 }
 
-func (c *Client) ListGCSPathDir(gcsDirPath string, maxResults *int64) ([]*storage.ObjectAttrs, error) {
+func (c *Client) ListGCSPathDir(gcsDirPath string, includeDirObjects bool, maxResults *int64) ([]*storage.ObjectAttrs, error) {
 	bucket, gcsDir, err := SplitGCSPath(gcsDirPath)
 	if err != nil {
 		return nil, err
 	}
-	return c.ListGCSDir(bucket, gcsDir, maxResults)
+	return c.ListGCSDir(bucket, gcsDir, includeDirObjects, maxResults)
 }
 
 // This behaves like you'd expect `ls` to behave on a local file system

pkg/operator/config/wrappers.go

@@ -159,7 +159,7 @@ func ListBucketPrefix(prefix string, maxResults *int64) ([]*storage.ObjectAttrs,
 		}
 		return nil, s3Objects, nil
 	case types.GCPProviderType:
-		gcsObjects, err := GCP.ListGCSPrefix(GCPCoreConfig.Bucket, prefix, maxResults)
+		gcsObjects, err := GCP.ListGCSPrefix(GCPCoreConfig.Bucket, prefix, false, maxResults)
 		if err != nil {
 			return nil, nil, err
 		}

pkg/types/spec/utils.go

@@ -153,7 +153,7 @@ func validateDirModels(
 			return nil, err
 		}
 
-		gcsObjects, err := gcpClient.ListGCSPathDir(modelPath, nil)
+		gcsObjects, err := gcpClient.ListGCSPathDir(modelPath, false, nil)
 		if err != nil {
 			return nil, err
 		}
@@ -287,7 +287,7 @@ func validateModels(
 			}
 			modelPrefix = s.EnsureSuffix(modelPrefix, "/")
 
-			gcsObjects, err := gcpClient.ListGCSPathDir(modelPath, nil)
+			gcsObjects, err := gcpClient.ListGCSPathDir(modelPath, false, nil)
 			if err != nil {
 				return nil, errors.Wrap(err, modelNameWrapStr)
 			}