Add fluid code-release changes in worldgen/tools/submit_asset_cache.py

worldgen/tools/submit_asset_cache.py

@@ -0,0 +1,58 @@
+# Copyright (c) Princeton University.
+# This source code is licensed under the BSD 3-clause license found in the LICENSE file in the root directory of this source tree.
+
+# Authors: Karhan Kayan
+
+import submitit
+import argparse
+from pathlib import Path
+import sys
+import os
+import time 
+
+sys.path.append(str(Path(os.path.split(os.path.abspath(__file__))[0]) / ".."))
+
+
+def get_slurm_banned_nodes(config_path=None):
+    if config_path is None:
+        return []
+    with Path(config_path).open('r') as f:
+        return list(f.read().split())
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('-b', '--blender_path', type=str)
+parser.add_argument('-f', '--asset_folder', type=str)
+parser.add_argument('-a', '--assets', nargs='+', default=[
+    'CachedBushFactory',
+    'CachedTreeFactory',
+    'CachedCactusFactory',
+    # 'CachedCreatureFactory',
+    'CachedBoulderFactory'
+])
+parser.add_argument('-n', '--number', type=int, default=1)
+parser.add_argument('-s', '--start_frame', type=int, default=-20)
+parser.add_argument('-d', '--simulation_duration', type=int, default=24*20+20)
+# parser.add_argument('-r', '--resolution', type=int)
+# parser.add_argument('--dissolve_speed', type=int, default=25)
+# parser.add_argument('--dom_scale', type=int, default=1)
+args = parser.parse_args()
+
+Path(args.asset_folder).mkdir(parents=True, exist_ok=True)
+
+
+for asset in args.assets:
+    for i in range(args.number):
+        cmd = f"{args.blender_path} --background -noaudio --python fluid/run_asset_cache.py -- -f {args.asset_folder}/ -a {asset} -s {args.start_frame} -d {args.simulation_duration}".split(" ")
+        print(cmd)
+        executor = submitit.AutoExecutor(folder=str(Path(args.asset_folder) / "logs"))
+        executor.update_parameters(
+            mem_gb=16,
+            name=f"{asset}_{i}",
+            cpus_per_task=4,
+            timeout_min=60*24,
+            slurm_account="pvl",
+            slurm_exclude= "node408,node409",
+        )
+        render_fn = submitit.helpers.CommandFunction(cmd)
+        executor.submit(render_fn)