Add status tracking (#72)

Author: tchaton

.github/workflows/ci-checks.yml

@@ -12,22 +12,21 @@ concurrency:
 
 jobs:
   check-precommit:
-    uses: Lightning-AI/utilities/.github/workflows/check-precommit.yml@v0.10.1
+    uses: Lightning-AI/utilities/.github/workflows/check-precommit.yml@main
 
   check-typing:
     # TODO: switch to main after fix lends
     uses: Lightning-AI/utilities/.github/workflows/check-typing.yml@main
     with:
       actions-ref: main
-      source-dir: ""
 
   check-schema:
-    uses: Lightning-AI/utilities/.github/workflows/check-schema.yml@v0.10.1
+    uses: Lightning-AI/utilities/.github/workflows/check-schema.yml@main
     with:
       azure-dir: ""
 
   check-package:
-    uses: Lightning-AI/utilities/.github/workflows/check-package.yml@v0.10.1
+    uses: Lightning-AI/utilities/.github/workflows/check-package.yml@main
     with:
       actions-ref: v0.10.1
       import-name: "litdata"

docs/source/conf.py

@@ -16,19 +16,22 @@
 
 _PATH_HERE = os.path.abspath(os.path.dirname(__file__))
 _PATH_ROOT = os.path.realpath(os.path.join(_PATH_HERE, "..", ".."))
+_PATH_SOURCE = os.path.join(_PATH_ROOT, "src")
 sys.path.insert(0, os.path.abspath(_PATH_ROOT))
 
 SPHINX_MOCK_REQUIREMENTS = int(os.environ.get("SPHINX_MOCK_REQUIREMENTS", True))
 
 # alternative https://stackoverflow.com/a/67692/4521646
-spec = spec_from_file_location("litdata/__about__.py", os.path.join(_PATH_ROOT, "litdata", "__about__.py"))
+spec = spec_from_file_location(
+    "litdata/__about__.py", os.path.join(_PATH_SOURCE, "litdata", "__about__.py")
+)
 about = module_from_spec(spec)
 spec.loader.exec_module(about)
 
 # -- Project information -----------------------------------------------------
 
 # this name shall match the project name in Github as it is used for linking to code
-project = "lit-data"
+project = "litdata"
 copyright = about.__copyright__
 author = about.__author__
 

pyproject.toml

@@ -88,7 +88,7 @@ lint.ignore-init-module-imports = true
     "S501",  # Probable use of `requests` call with `verify=False` disabling SSL certificate checks
     "S108",  # Probable insecure usage of temporary file or directory: "/tmp/data/MNIST"
 ]
-"litdata/**" = [
+"src/**" = [
     "S101",  # todo: Use of `assert` detected
     "S105", "S106", "S107",  # todo: Possible hardcoded password: ...
     "S113",  # todo: Probable use of requests call without timeout

setup.py

@@ -5,14 +5,15 @@
 from pathlib import Path
 
 from pkg_resources import parse_requirements
-from setuptools import setup
+from setuptools import find_packages, setup
 
 _PATH_ROOT = os.path.dirname(__file__)
-_PATH_REQUIRES = os.path.join(_PATH_ROOT, "requirements")
+_PATH_SOURCE = os.path.join(_PATH_ROOT, "src")
+_PATH_REQUIRES = os.path.join(_PATH_ROOT, "_requirements")
 
 
 def _load_py_module(fname, pkg="litdata"):
-    spec = spec_from_file_location(os.path.join(pkg, fname), os.path.join(_PATH_ROOT, pkg, fname))
+    spec = spec_from_file_location(os.path.join(pkg, fname), os.path.join(_PATH_SOURCE, pkg, fname))
     py = module_from_spec(spec)
     spec.loader.exec_module(py)
     return py
@@ -60,6 +61,8 @@ def _prepare_extras(requirements_dir: str = _PATH_REQUIRES, skip_files: tuple =
     download_url="https://github.com/Lightning-AI/litdata",
     license=about.__license__,
     long_description=readme,
+    packages=find_packages(where="src"),
+    package_dir={"": "src"},
     long_description_content_type="text/markdown",
     include_package_data=True,
     zip_safe=False,

litdata/CHANGELOG.md renamed to src/litdata/CHANGELOG.md

@@ -953,6 +953,9 @@ def run(self, data_recipe: DataRecipe) -> None:
             leave=True,
             dynamic_ncols=True,
         )
+        num_nodes = _get_num_nodes()
+        node_rank = _get_node_rank()
+        total_num_items = len(user_items)
 
         while True:
             try:
@@ -973,6 +976,10 @@ def run(self, data_recipe: DataRecipe) -> None:
             if current_total == num_items:
                 break
 
+            if _IS_IN_STUDIO and node_rank == 0:
+                with open("status.json", "w") as f:
+                    json.dump({"progress": str(100 * current_total * num_nodes / total_num_items) + "%"}, f)
+
             # Exit early if all the workers are done.
             # This means there were some kinda of errors.
             if all(not w.is_alive() for w in self.workers):
@@ -981,8 +988,6 @@ def run(self, data_recipe: DataRecipe) -> None:
 
         pbar.close()
 
-        num_nodes = _get_num_nodes()
-        node_rank = _get_node_rank()
         # TODO: Understand why it hangs.
         if num_nodes == 1:
             for w in self.workers: