minor tweaks to skypilot setup (#348)

* Fixes numba bug with print being defined inside a function rather than
a module.
* Fixes sqlalchemy required by skypilot
* Some docs and convenience updates

Author: jder

CONTRIBUTING.md

@@ -170,9 +170,10 @@ You can run `uv run -m ocean_emulators.train --help` to see all the options avai
 To learn more about other datasets used during training, please see the _Data Engineering_ section below.
 
 To run a remote training job with Skypilot, use the following command:
+
 ```shell
 # export WANDB_API_KEY=<my-key>  # Get your key at https://wandb.ai/authorize
-# sky launch -c fomo-cluster train.sky.yaml  --env WANDB_API_KEY --env-file <my-vars>.env
+uv run sky launch train.sky.yaml  --env WANDB_API_KEY --env-file <my-vars>.env
 ```
 
 Please read the docstring in the `train.sky.yaml` for more information.

pyproject.toml

@@ -38,6 +38,7 @@ dependencies = [
   "s3fs>=2025.5.1",
   "scipy>=1.15.2",
   "skypilot[aws,lambda]>=0.10",
+  "sqlalchemy>=2",                # skypilot requires sqlalchemy 2.0 but doesn't include that requirement
   "torch~=2.2.1",
   "torchinfo>=1.8",
   "wandb>=0.19.8",

src/ocean_emulators/utils/distributed.py

@@ -25,19 +25,22 @@ def set_seed(seed):
     # torch.use_deterministic_algorithms(True)
 
 
-def suppress_prints(is_master):
-    """This function disables printing when not in master process."""
-    builtin_print = builtins.print
+builtin_print = builtins.print
+
+
+def print(*args, **kwargs):
+    force = kwargs.pop("force", False)
+    force = force or (get_world_size() > 8)
+    if force:
+        now = datetime.datetime.now().time()
+        builtin_print(f"[{now}] ", end="")  # print with time stamp
+        builtin_print(*args, **kwargs)
 
-    def print(*args, **kwargs):
-        force = kwargs.pop("force", False)
-        force = force or (get_world_size() > 8)
-        if is_master or force:
-            now = datetime.datetime.now().time()
-            builtin_print(f"[{now}] ", end="")  # print with time stamp
-            builtin_print(*args, **kwargs)
 
-    builtins.print = print
+def suppress_prints(is_master):
+    """This function disables printing when not in master process."""
+    if not is_master:
+        builtins.print = print
 
 
 def suppress_logging(is_master):

train.sky.yaml

@@ -6,11 +6,13 @@
 # ```
 # sky check -v
 # ```
+# You will generally need to create a dedicated IAM user for skypilot, and put
+# your lambda credentials in ~/.lambda_cloud/ -- the `sky check -v` command has more details.
 #
 # To create or re-use a cluster:
 # ```
 # export WANDB_API_KEY=<my-key>  # Get your key at https://wandb.ai/authorize
-# sky launch -c fomo-cluster train.sky.yaml --env-file <my-env>.env --env WANDB_API_KEY
+# uv run sky launch train.sky.yaml --env-file <my-env>.env --env WANDB_API_KEY
 # ```
 # As the job progresses, skypilot will provide commands to locally monitor each stage.
 #
@@ -44,6 +46,7 @@ envs:
   # We assume that data in this directory is named in the same pattern as
   # the $DATA_CONFIG yaml.
   DATA_ROOT_PATH: null # data/om4_onedeg_compact  # NOTE: path cannot end in a slash `/`.
+  ARGS: ""
 
 file_mounts:
   # TODO(alxmrs): Consider adding a data/ COPY file_mount here once `mount_options` feature exists.
@@ -56,6 +59,8 @@ file_mounts:
 workdir: .
 
 setup: |
+  set -ex
+
   mkdir ~/data/
   rclone config create s3-source s3 provider=AWS env_auth=true
 
@@ -64,6 +69,8 @@ setup: |
   rclone copy --progress s3-source:${DATA_BUCKET}/${DATA_ROOT_PATH}/OM4.zarr ~/data/OM4.zarr --transfers=32 --ignore-existing
 
 run: |
+  set -ex
+
   MASTER_ADDR=$(echo "$SKYPILOT_NODE_IPS" | head -n1)
   echo "Starting distributed training, head node: $MASTER_ADDR"
 
@@ -86,4 +93,4 @@ run: |
     --experiment.data_root=${DATA_ROOT} \
     --experiment.base_output_dir=/outputs/  \
     --experiment.wandb.mode=online \
-    --samudra.checkpointing=all
+    ${ARGS}