DaCe VRAM pooling

constraints.txt

@@ -1,5 +1,5 @@
 #
-# This file is autogenerated by pip-compile
+# This file is autogenerated by pip-compile with python 3.8
 # To update, run:
 #
 #    pip-compile --output-file=constraints.txt driver/requirements.txt dsl/requirements.txt external/gt4py/setup.cfg fv3core/requirements.txt fv3gfs-physics/requirements.txt pace-util/requirements.txt requirements_dev.txt requirements_docs.txt requirements_lint.txt
@@ -34,7 +34,7 @@ attrs==21.2.0
     #   pytest
 babel==2.9.1
     # via sphinx
-backports.entry-points-selectable==1.1.1
+backports-entry-points-selectable==1.1.1
     # via virtualenv
 black==22.3.0
     # via
@@ -76,8 +76,6 @@ click==8.0.1
     #   pip-tools
 cloudpickle==2.0.0
     # via dask
-cmake==3.22.4
-    # via dace
 commonmark==0.9.1
     # via recommonmark
 coverage==5.5
@@ -88,6 +86,13 @@ cytoolz==0.11.2
     # via
     #   gt4py
     #   gt4py (external/gt4py/setup.cfg)
+dace==0.14
+    # via
+    #   -r driver/requirements.txt
+    #   -r dsl/requirements.txt
+    #   -r fv3core/requirements/requirements_dace.txt
+    #   -r requirements_dev.txt
+    #   pace-dsl
 dacite==1.6.0
     # via
     #   -r driver/requirements.txt
@@ -105,8 +110,6 @@ dill==0.3.5.1
     # via dace
 distlib==0.3.2
     # via virtualenv
-distro==1.7.0
-    # via scikit-build
 docutils==0.16
     # via
     #   recommonmark
@@ -163,15 +166,15 @@ google-api-core==2.0.0
     # via
     #   google-cloud-core
     #   google-cloud-storage
-google-auth-oauthlib==0.4.5
-    # via gcsfs
 google-auth==2.0.1
     # via
     #   gcsfs
     #   google-api-core
     #   google-auth-oauthlib
     #   google-cloud-core
     #   google-cloud-storage
+google-auth-oauthlib==0.4.5
+    # via gcsfs
 google-cloud-core==2.0.0
     # via google-cloud-storage
 google-cloud-storage==1.42.0
@@ -238,15 +241,15 @@ multidict==5.1.0
     # via
     #   aiohttp
     #   yarl
+mypy==0.790
+    # via
+    #   -r fv3gfs-physics/requirements.txt
+    #   -r pace-util/requirements.txt
 mypy-extensions==0.4.3
     # via
     #   black
     #   mypy
     #   typing-inspect
-mypy==0.790
-    # via
-    #   -r fv3gfs-physics/requirements.txt
-    #   -r pace-util/requirements.txt
 netcdf4==1.5.7
     # via
     #   -r driver/requirements.txt
@@ -292,7 +295,6 @@ packaging==21.0
     #   gt4py
     #   gt4py (external/gt4py/setup.cfg)
     #   pytest
-    #   scikit-build
     #   sphinx
     #   tox
 pandas==1.3.2
@@ -326,12 +328,12 @@ py==1.10.0
     #   pytest
     #   pytest-forked
     #   tox
-pyasn1-modules==0.2.8
-    # via google-auth
 pyasn1==0.4.8
     # via
     #   pyasn1-modules
     #   rsa
+pyasn1-modules==0.2.8
+    # via google-auth
 pybind11==2.8.1
     # via
     #   gt4py
@@ -350,6 +352,21 @@ pygments==2.10.0
     # via sphinx
 pyparsing==2.4.7
     # via packaging
+pytest==6.2.4
+    # via
+    #   -r driver/requirements.txt
+    #   -r fv3core/requirements/requirements_base.txt
+    #   -r requirements_dev.txt
+    #   pytest-cache
+    #   pytest-cov
+    #   pytest-datadir
+    #   pytest-dependency
+    #   pytest-factoryboy
+    #   pytest-forked
+    #   pytest-profiling
+    #   pytest-regressions
+    #   pytest-subtests
+    #   pytest-xdist
 pytest-cache==1.0
     # via -r fv3core/requirements/requirements_base.txt
 pytest-cov==2.12.1
@@ -378,21 +395,6 @@ pytest-subtests==0.5.0
     #   -r requirements_dev.txt
 pytest-xdist==2.3.0
     # via -r fv3core/requirements/requirements_base.txt
-pytest==6.2.4
-    # via
-    #   -r driver/requirements.txt
-    #   -r fv3core/requirements/requirements_base.txt
-    #   -r requirements_dev.txt
-    #   pytest-cache
-    #   pytest-cov
-    #   pytest-datadir
-    #   pytest-dependency
-    #   pytest-factoryboy
-    #   pytest-forked
-    #   pytest-profiling
-    #   pytest-regressions
-    #   pytest-subtests
-    #   pytest-xdist
 python-dateutil==2.8.2
     # via
     #   faker
@@ -411,8 +413,6 @@ pyyaml==5.4.1
     #   pytest-regressions
 recommonmark==0.7.1
     # via -r requirements_docs.txt
-requests-oauthlib==1.3.0
-    # via google-auth-oauthlib
 requests==2.26.0
     # via
     #   dace
@@ -421,10 +421,10 @@ requests==2.26.0
     #   google-cloud-storage
     #   requests-oauthlib
     #   sphinx
+requests-oauthlib==1.3.0
+    # via google-auth-oauthlib
 rsa==4.7.2
     # via google-auth
-scikit-build==0.15.0
-    # via dace
 scipy==1.7.1
     # via
     #   -r fv3core/requirements/requirements_base.txt
@@ -447,6 +447,13 @@ snowballstemmer==2.1.0
     # via sphinx
 sortedcontainers==2.4.0
     # via hypothesis
+sphinx==4.1.2
+    # via
+    #   -r requirements_docs.txt
+    #   recommonmark
+    #   sphinx-argparse
+    #   sphinx-gallery
+    #   sphinx-rtd-theme
 sphinx-argparse==0.3.1
     # via -r requirements_docs.txt
 sphinx-gallery==0.10.1
@@ -455,13 +462,6 @@ sphinx-rtd-theme==0.5.2
     # via
     #   -r pace-util/requirements.txt
     #   -r requirements_docs.txt
-sphinx==4.1.2
-    # via
-    #   -r requirements_docs.txt
-    #   recommonmark
-    #   sphinx-argparse
-    #   sphinx-gallery
-    #   sphinx-rtd-theme
 sphinxcontrib-applehelp==1.0.2
     # via sphinx
 sphinxcontrib-devhelp==1.0.2
@@ -533,7 +533,6 @@ wheel==0.37.0
     #   -r pace-util/requirements.txt
     #   astunparse
     #   pip-tools
-    #   scikit-build
 xarray==0.19.0
     # via
     #   -r driver/requirements.txt

driver/pace/driver/tools.py

@@ -0,0 +1,35 @@
+import os
+from typing import Optional
+
+import click
+
+from pace.dsl.dace.utils import count_memory_from_path
+
+
+# Count the memory from a given SDFG
+ACTION_SDFG_MEMORY_COUNT = "sdfg_memory_count"
+
+
+@click.command()
+@click.argument(
+    "action",
+    required=True,
+    type=click.Choice([ACTION_SDFG_MEMORY_COUNT]),
+)
+@click.option(
+    "--sdfg_path",
+    type=click.STRING,
+)
+@click.option("--report_detail", is_flag=True, type=click.BOOL, default=False)
+def command_line(action: str, sdfg_path: Optional[str], report_detail: Optional[bool]):
+    """
+    Run tooling.
+    """
+    if action == ACTION_SDFG_MEMORY_COUNT:
+        if sdfg_path is None or not os.path.exists(sdfg_path):
+            raise RuntimeError(f"Can't load SDFG {sdfg_path}")
+        print(count_memory_from_path(sdfg_path, detail_report=report_detail))
+
+
+if __name__ == "__main__":
+    command_line()

driver/requirements.txt

@@ -7,4 +7,4 @@ numpy
 netCDF4
 xarray
 zarr
-git+https://github.com/spcl/dace[email protected]
+dace>=0.14

dsl/pace/dsl/dace/build.py

@@ -110,6 +110,7 @@ def get_sdfg_path(
     """
     import os
 
+    # TODO: check DaceConfig for cache.strategy == name
     # Guarding against bad usage of this function
     if config.get_orchestrate() != DaCeOrchestration.Run:
         return None

dsl/pace/dsl/dace/dace_config.py

@@ -7,6 +7,9 @@
 from pace.util.communicator import CubedSphereCommunicator
 
 
+TEMPORARY_DEACTIVATE_DISTRIBUTED_DACE_COMPILE = True
+
+
 class DaCeOrchestration(enum.Enum):
     """
     Orchestration mode for DaCe
@@ -139,7 +142,12 @@ def __init__(
         if communicator:
             self.my_rank = communicator.rank
             self.rank_size = communicator.comm.Get_size()
-            self.target_rank = get_target_rank(self.my_rank, communicator.partitioner)
+            if TEMPORARY_DEACTIVATE_DISTRIBUTED_DACE_COMPILE:
+                self.target_rank = communicator.rank
+            else:
+                self.target_rank = get_target_rank(
+                    self.my_rank, communicator.partitioner
+                )
             self.layout = communicator.partitioner.layout
         else:
             self.my_rank = 0
@@ -170,6 +178,9 @@ def get_backend(self) -> str:
     def get_orchestrate(self) -> DaCeOrchestration:
         return self._orchestrate
 
+    def sync_debug(self) -> bool:
+        return dace.config.Config.get("compiler", "cuda", "syncdebug")
+
     def as_dict(self) -> Dict[str, Any]:
         return {
             "_orchestrate": str(self._orchestrate.name),