eotdl 2023.10.25.post10__py3-none-any.whl → 2023.11.2.post2__py3-none-any.whl

eotdl/datasets/ingest.py

@@ -1,13 +1,10 @@
 from pathlib import Path
-from glob import glob
 import yaml
-from tqdm import tqdm
-import os
 
 from ..auth import with_auth
 from .metadata import Metadata
-from ..repos import DatasetsAPIRepo, FilesAPIRepo
-from .utils import calculate_checksum
+from ..repos import DatasetsAPIRepo
+from ..files import ingest_files
 
 
 def ingest_dataset(path, verbose=False, logger=print):
@@ -19,160 +16,34 @@ def ingest_dataset(path, verbose=False, logger=print):
     return ingest_folder(path, verbose, logger)
 
 
+def retrieve_dataset(metadata, user):
+    repo = DatasetsAPIRepo()
+    data, error = repo.retrieve_dataset(metadata.name)
+    # print(data, error)
+    if data and data["uid"] != user["sub"]:
+        raise Exception("Dataset already exists.")
+    if error and error == "Dataset doesn't exist":
+        # create dataset
+        data, error = repo.create_dataset(metadata.dict(), user["id_token"])
+        # print(data, error)
+        if error:
+            raise Exception(error)
+        data["id"] = data["dataset_id"]
+    return data["id"]
+
+
 @with_auth
 def ingest_folder(folder, verbose=False, logger=print, user=None):
-    repo, files_repo = DatasetsAPIRepo(), FilesAPIRepo()
-    logger(f"Uploading directory {folder}...")
-    # get all files in directory recursively
-    items = [Path(item) for item in glob(str(folder) + "/**/*", recursive=True)]
-    # remove directories
-    items = [item for item in items if not item.is_dir()]
-    if len(items) == 0:
-        raise Exception("No files found in directory")
-    if not any(item.name == "metadata.yml" for item in items):
-        raise Exception("metadata.yml not found in directory")
+    repo = DatasetsAPIRepo()
     # load metadata
     metadata = yaml.safe_load(open(folder.joinpath("metadata.yml"), "r").read()) or {}
     metadata = Metadata(**metadata)
-    # remove metadata.yml from files
-    items = [item for item in items if item.name != "metadata.yml"]
-    # if zip or tar file, send error
-    if any(item.suffix.endswith((".zip", ".tar", ".tar.gz", ".gz")) for item in items):
-        raise Exception(
-            f"At least one zip, tar or gz file found in {folder}, please unzip and try again"
-        )
-    # create dataset
-    data, error = repo.create_dataset(metadata.dict(), user["id_token"])
-    # dataset may already exist, and will return an error, but if user is owner continue ingesting files
-    current_files = []
-    if error:
-        data, error2 = repo.retrieve_dataset(metadata.name)
-        if error2:
-            raise Exception(error)
-        if data["uid"] != user["sub"]:
-            raise Exception("Dataset already exists.")
-        data["dataset_id"] = data["id"]
-    dataset_id = data["dataset_id"]
-    # create new version
-    data, error = repo.create_version(dataset_id, user["id_token"])
-    if error:
-        raise Exception(error)
-    version = data["version"]
-    # upload files
-    current_files = []
-    if version > 1:
-        current_files, error = files_repo.retrieve_dataset_files(
-            dataset_id, version - 1
-        )
-        if error:
-            # print("retreive dataset files error: ", error)
-            current_files = []
-    for item in tqdm(items, desc="Uploading files", unit="files", disable=verbose):
-        data = ingest_file(
-            str(item),
-            dataset_id,
-            version,
-            str(item.relative_to(folder).parent),
-            logger=logger,
-            verbose=verbose,
-            user=user,
-            current_files=current_files,
-        )
-    return data
-
-
-def ingest_file(
-    file,
-    dataset_id,
-    version,
-    parent,
-    logger=None,
-    verbose=True,
-    root=None,
-    user=None,
-    current_files=[],
-):
-    id_token = user["id_token"]
-    if verbose:
-        logger(f"Uploading file {file}...")
-    repo = FilesAPIRepo()
-    if file.startswith("http://") or file.startswith("https://"):
-        raise NotImplementedError("URL ingestion not implemented yet")
-        # data, error = repo.ingest_file_url(file, dataset_id, id_token)
-    else:
-        file_path = Path(file)
-        if not file_path.is_absolute():
-            # file_path = glob(
-            #     str(root) + "/**/" + os.path.basename(file_path),
-            #     recursive=True,
-            # )
-            # if len(file_path) == 0:
-            #     raise Exception(f"File {file} not found")
-            # elif len(file_path) > 1:
-            #     raise Exception(f"Multiple files found for {file}")
-            # file_path = file_path[0]
-            file_path = str(file_path.absolute())
-        if verbose:
-            logger("Computing checksum...")
-        checksum = calculate_checksum(file_path)
-        # check if file already exists in dataset
-        filename = os.path.basename(file_path)
-        if parent != ".":
-            filename = parent + "/" + filename
-        if len(current_files) > 0:
-            matches = [
-                f
-                for f in current_files
-                if f["filename"] == filename and f["checksum"] == checksum
-            ]  # this could slow down ingestion in large datasets... should think of faster search algos, puede que sea mejor hacer el re-upload simplemente...
-            if len(matches) == 1:
-                if verbose:
-                    print(f"File {file_path} already exists in dataset, skipping...")
-                data, error = repo.ingest_existing_file(
-                    filename,
-                    dataset_id,
-                    version,
-                    matches[0]["version"],
-                    id_token,
-                    checksum,
-                )
-                if error:
-                    raise Exception(error)
-                if verbose:
-                    logger("Done")
-                return data
-        if verbose:
-            logger("Ingesting file...")
-        filesize = os.path.getsize(file_path)
-        # ingest small file
-        if filesize < 1024 * 1024 * 16:  # 16 MB
-            data, error = repo.ingest_file(
-                file_path,
-                dataset_id,
-                version,
-                parent,
-                id_token,
-                checksum,
-            )
-            if error:
-                raise Exception(error)
-            if verbose:
-                logger("Done")
-            return data
-        raise NotImplementedError("Large file ingestion not implemented yet")
-        # # ingest large file
-        # upload_id, parts = repo.prepare_large_upload(
-        #     file_path, dataset_id, checksum, id_token
-        # )
-        # repo.ingest_large_dataset(file_path, upload_id, id_token, parts)
-        # if verbose:
-        #     logger("\nCompleting upload...")
-        # data, error = repo.complete_upload(id_token, upload_id)
-    if error:
-        raise Exception(error)
-    if verbose:
-        logger("Done")
-    return data
+    # retrieve dataset (create if doesn't exist)
+    dataset_id = retrieve_dataset(metadata, user)
+    # ingest files
+    return ingest_files(
+        repo, dataset_id, folder, verbose, logger, user, endpoint="datasets"
+    )
 
 
 # @with_auth

eotdl/datasets/retrieve.py

@@ -20,7 +20,7 @@ def retrieve_dataset(name):
 
 def retrieve_dataset_files(dataset_id, version):
     repo = FilesAPIRepo()
-    data, error = repo.retrieve_dataset_files(dataset_id, version)
+    data, error = repo.retrieve_files(dataset_id, "datasets", version)
     if error:
         raise Exception(error)
     return data

eotdl/files/__init__.py

@@ -0,0 +1 @@
+from .ingest import ingest_files

eotdl/files/ingest.py

@@ -0,0 +1,175 @@
+from pathlib import Path
+import os
+from tqdm import tqdm
+import zipfile
+import io
+from glob import glob
+import os
+
+from ..repos import FilesAPIRepo
+from ..shared import calculate_checksum
+from ..shared import calculate_checksum
+
+
+def retrieve_files(folder):
+    # get all files in directory recursively
+    items = [Path(item) for item in glob(str(folder) + "/**/*", recursive=True)]
+    if not any(item.name == "metadata.yml" for item in items):
+        raise Exception("metadata.yml not found in directory")
+    # remove directories
+    items = [item for item in items if not item.is_dir()]
+    if len(items) == 0:
+        raise Exception("No files found in directory")
+    return items
+
+
+def prepare_item(item, folder):
+    return {
+        "filename": item.name,
+        "path": str(item.relative_to(folder)),
+        "absolute_path": item.absolute(),
+        "size": os.path.getsize(item.absolute()),
+        "checksum": calculate_checksum(item.absolute()),
+    }
+
+
+def generate_batches(files, max_batch_size=1024 * 1024 * 10, max_batch_files=10):
+    batches = []
+    for item in tqdm(files):
+        if not batches:
+            batches.append([item])
+            continue
+        if max_batch_size:
+            size_check = sum([i["size"] for i in batches[-1]]) < max_batch_size
+        else:
+            size_check = True
+        if size_check and len(batches[-1]) < max_batch_files:
+            batches[-1].append(item)
+        else:
+            batches.append([item])
+    return batches
+
+
+def compress_batch(batch):
+    memory_file = io.BytesIO()
+    with zipfile.ZipFile(memory_file, "w") as zf:
+        for f in batch:
+            zf.write(f["absolute_path"], arcname=f["path"])
+    memory_file.seek(0)
+    return memory_file
+
+
+def generate_files_lists(
+    items, folder, dataset_or_model_id, endpoint, logger, max_size=1024 * 1024 * 16
+):
+    files_repo = FilesAPIRepo()
+    current_files, error = files_repo.retrieve_files(
+        dataset_or_model_id, "models", endpoint
+    )
+    # print(len(current_files), len(items) - len(current_files))
+    # print(current_files, error)
+    if error:
+        current_files = []
+    # generate list of files to upload
+    logger("generating list of files to upload...")
+    upload_files, existing_files, large_files = [], [], []
+    current_names = [f["filename"] for f in current_files]
+    current_checksums = [f["checksum"] for f in current_files]
+    for item in tqdm(items):
+        data = prepare_item(item, folder)
+        if data["path"] in current_names and data["checksum"] in current_checksums:
+            existing_files.append(data)
+        else:
+            if data["size"] > max_size:
+                large_files.append(data)
+            else:
+                upload_files.append(data)
+    if len(upload_files) == 0 and len(large_files) == 0:
+        raise Exception("No files to upload")
+    return upload_files, existing_files, large_files
+
+
+def create_new_version(repo, dataset_or_model_id, user):
+    data, error = repo.create_version(dataset_or_model_id, user["id_token"])
+    if error:
+        raise Exception(error)
+    return data["version"]
+
+
+def ingest_files(repo, dataset_or_model_id, folder, verbose, logger, user, endpoint):
+    files_repo = FilesAPIRepo()
+    logger(f"Uploading directory {folder}...")
+    items = retrieve_files(folder)
+    # retrieve files
+    upload_files, existing_files, large_files = generate_files_lists(
+        items, folder, dataset_or_model_id, endpoint, logger
+    )
+    logger(f"{len(upload_files) + len(large_files)} new files will be ingested")
+    logger(f"{len(existing_files)} files already exist in dataset")
+    logger(f"{len(large_files)} large files will be ingested separately")
+    # create new version
+    version = create_new_version(repo, dataset_or_model_id, user)
+    logger("New version created, version: " + str(version))
+    # ingest new large files
+    if len(large_files) > 0:
+        logger("ingesting large files...")
+        for file in large_files:
+            logger("ingesting file: " + file["path"])
+            upload_id, parts = files_repo.prepare_large_upload(
+                file["path"],
+                dataset_or_model_id,
+                file["checksum"],
+                user["id_token"],
+                endpoint,
+            )
+            # print(upload_id, parts)
+            files_repo.ingest_large_file(
+                file["absolute_path"],
+                file["size"],
+                upload_id,
+                user["id_token"],
+                parts,
+                endpoint,
+            )
+            files_repo.complete_upload(user["id_token"], upload_id, version, endpoint)
+    # ingest new small files in batches
+    if len(upload_files) > 0:
+        logger("generating batches...")
+        batches = generate_batches(upload_files)
+        logger(
+            f"Uploading {len(upload_files)} small files in {len(batches)} batches..."
+        )
+        repo = FilesAPIRepo()
+        for batch in tqdm(
+            batches, desc="Uploading batches", unit="batches", disable=verbose
+        ):
+            # compress batch
+            memory_file = compress_batch(batch)
+            # ingest batch
+            data, error = repo.ingest_files_batch(
+                memory_file,
+                [f["checksum"] for f in batch],
+                dataset_or_model_id,
+                user["id_token"],
+                endpoint,
+                version,
+            )
+    # ingest existing files
+    if len(existing_files) > 0:
+        batches = generate_batches(existing_files, max_batch_size=None)
+        for batch in tqdm(
+            batches,
+            desc="Ingesting existing files",
+            unit="batches",
+            disable=verbose,
+        ):
+            data, error = files_repo.add_files_batch_to_version(
+                batch,
+                dataset_or_model_id,
+                version,
+                user["id_token"],
+                endpoint,
+            )
+            if error:
+                raise Exception(error)
+    return data

eotdl/models/__init__.py

@@ -0,0 +1,3 @@
+from .retrieve import retrieve_models
+from .ingest import ingest_model
+from .download import download_model

eotdl/models/download.py

@@ -0,0 +1,119 @@
+import os
+from pathlib import Path
+from tqdm import tqdm
+
+from ..auth import with_auth
+from .retrieve import retrieve_model, retrieve_model_files
+from ..shared import calculate_checksum
+from ..repos import FilesAPIRepo
+
+
+@with_auth
+def download_model(
+    model_name,
+    version=None,
+    path=None,
+    logger=None,
+    assets=False,
+    force=False,
+    verbose=False,
+    user=None,
+    file=None,
+):
+    model = retrieve_model(model_name)
+    if version is None:
+        version = sorted(model["versions"], key=lambda v: v["version_id"])[-1][
+            "version_id"
+        ]
+    else:
+        assert version in [
+            v["version_id"] for v in model["versions"]
+        ], f"Version {version} not found"
+    download_base_path = os.getenv(
+        "EOTDL_DOWNLOAD_PATH", str(Path.home()) + "/.cache/eotdl/models"
+    )
+    if path is None:
+        download_path = download_base_path + "/" + model_name + "/v" + str(version)
+    else:
+        download_path = path + "/" + model_name + "/v" + str(version)
+    # check if model already exists
+    if os.path.exists(download_path) and not force:
+        os.makedirs(download_path, exist_ok=True)
+        raise Exception(
+            f"model `{model['name']} v{str(version)}` already exists at {download_path}. To force download, use force=True or -f in the CLI."
+        )
+    if model["quality"] == 0:
+        if file:
+            raise NotImplementedError("Downloading a specific file is not implemented")
+            # files = [f for f in model["files"] if f["name"] == file]
+            # if not files:
+            #     raise Exception(f"File {file} not found")
+            # if len(files) > 1:
+            #     raise Exception(f"Multiple files with name {file} found")
+            # dst_path = download(
+            #     model,
+            #     model["id"],
+            #     file,
+            #     files[0]["checksum"],
+            #     download_path,
+            #     user,
+            # )
+            # return Outputs(dst_path=dst_path)
+        model_files = retrieve_model_files(model["id"], version)
+        repo = FilesAPIRepo()
+        for file in tqdm(model_files, disable=verbose, unit="file"):
+            filename, file_version = file["filename"], file["version"]
+            if verbose:
+                logger(f"Downloading {file['filename']}...")
+            dst_path = repo.download_file(
+                model["id"],
+                filename,
+                user["id_token"],
+                download_path,
+                file_version,
+                endpoint="models",
+            )
+            # if calculate_checksum(dst_path) != checksum:
+            #     logger(f"Checksum for {file} does not match")
+            if verbose:
+                logger(f"Done")
+        return "/".join(dst_path.split("/")[:-1])
+    else:
+        raise NotImplementedError("Downloading a STAC model is not implemented")
+    #     logger("Downloading STAC metadata...")
+    #     gdf, error = repo.download_stac(
+    #         model["id"],
+    #         user["id_token"],
+    #     )
+    #     if error:
+    #         raise Exception(error)
+    #     df = STACDataFrame(gdf)
+    #     # df.geometry = df.geometry.apply(lambda x: Polygon() if x is None else x)
+    #     path = path
+    #     if path is None:
+    #         path = download_base_path + "/" + model["name"]
+    #     df.to_stac(path)
+    #     # download assets
+    #     if assets:
+    #         logger("Downloading assets...")
+    #         df = df.dropna(subset=["assets"])
+    #         for row in tqdm(df.iterrows(), total=len(df)):
+    #             id = row[1]["stac_id"]
+    #             # print(row[1]["links"])
+    #             for k, v in row[1]["assets"].items():
+    #                 href = v["href"]
+    #                 repo.download_file_url(
+    #                     href, f"{path}/assets/{id}", user["id_token"]
+    #                 )
+    #     else:
+    #         logger("To download assets, set assets=True or -a in the CLI.")
+    #     return Outputs(dst_path=path)
+
+
+# @with_auth
+# def download_file_url(url, path, progress=True, logger=None, user=None):
+#     api_repo = APIRepo()
+#     download = DownloadFileURL(api_repo, logger, progress)
+#     inputs = DownloadFileURL.Inputs(url=url, path=path, user=user)
+#     outputs = download(inputs)
+#     return outputs.dst_path

eotdl/models/ingest.py

@@ -0,0 +1,47 @@
+from pathlib import Path
+import yaml
+
+from ..auth import with_auth
+from .metadata import Metadata
+from ..repos import ModelsAPIRepo
+from ..shared import calculate_checksum
+from ..files import ingest_files
+
+
+def ingest_model(path, verbose=False, logger=print):
+    path = Path(path)
+    if not path.is_dir():
+        raise Exception("Path must be a folder")
+    # if "catalog.json" in [f.name for f in path.iterdir()]:
+    #     return ingest_stac(path / "catalog.json", logger)
+    return ingest_folder(path, verbose, logger)
+
+
+def retrieve_model(metadata, user):
+    repo = ModelsAPIRepo()
+    data, error = repo.retrieve_model(metadata.name)
+    # print(data, error)
+    if data and data["uid"] != user["sub"]:
+        raise Exception("Model already exists.")
+    if error and error == "Model doesn't exist":
+        # create dataset
+        data, error = repo.create_model(metadata.dict(), user["id_token"])
+        # print(data, error)
+        if error:
+            raise Exception(error)
+        data["id"] = data["model_id"]
+    return data["id"]
+
+
+@with_auth
+def ingest_folder(folder, verbose=False, logger=print, user=None):
+    repo = ModelsAPIRepo()
+    # load metadata
+    metadata = yaml.safe_load(open(folder.joinpath("metadata.yml"), "r").read()) or {}
+    metadata = Metadata(**metadata)
+    # retrieve model (create if doesn't exist)
+    model_id = retrieve_model(metadata, user)
+    # ingest files
+    return ingest_files(
+        repo, model_id, folder, verbose, logger, user, endpoint="models"
+    )

eotdl/models/metadata.py

@@ -0,0 +1,16 @@
+from pydantic import BaseModel, validator
+from typing import List
+
+
+class Metadata(BaseModel):
+    authors: List[str]
+    license: str
+    source: str
+    name: str
+
+    # validate source is a URL
+    @validator("source")
+    def source_is_url(cls, v):
+        if not v.startswith("http") and not v.startswith("https"):
+            raise ValueError("source must be a URL")
+        return v

eotdl/models/retrieve.py

@@ -0,0 +1,26 @@
+from ..repos import ModelsAPIRepo, FilesAPIRepo
+
+
+def retrieve_models(name=None, limit=None):
+    api_repo = ModelsAPIRepo()
+    data, error = api_repo.retrieve_models(name, limit)
+    if data and not error:
+        models = [d["name"] for d in data] if data else []
+        return models
+    return []
+
+
+def retrieve_model(name):
+    repo = ModelsAPIRepo()
+    data, error = repo.retrieve_model(name)
+    if error:
+        raise Exception(error)
+    return data
+
+
+def retrieve_model_files(model_id, version):
+    repo = FilesAPIRepo()
+    data, error = repo.retrieve_files(model_id, "models", version)
+    if error:
+        raise Exception(error)
+    return data