Feature: Resolve recipes correctly before launching (#1529)

* fix to work with launcher recipes

* fix suffix for temp file

* fix path and error message

* fix for recipes from launcher

* resolve recipes correctly

* fix imports

* reformat message to avoid code-doc test issue

* code style fix

* code style fix

* code style fix

* code style fix

* code style fix

* code style fix

* code style fix

* code style fix

* code style fix

* doc formatting

* check if resolver exists before registering

Author: schinmayee

src/sagemaker/pytorch/estimator.py

@@ -14,13 +14,15 @@
 from __future__ import absolute_import
 
 import logging
+import math
 import os
 import shutil
 import tempfile
 from typing import Union, Optional, Dict
 from urllib.request import urlretrieve
 
-from omegaconf import OmegaConf
+import omegaconf
+from omegaconf import OmegaConf, dictconfig
 from packaging.version import Version
 
 from sagemaker.estimator import Framework, EstimatorBase
@@ -42,6 +44,19 @@
 logger = logging.getLogger("sagemaker")
 
 
+def _try_resolve_recipe(recipe, key=None):
+    """Try to resolve recipe and return resolved recipe."""
+    if key is not None:
+        recipe = dictconfig.DictConfig({key: recipe})
+    try:
+        OmegaConf.resolve(recipe)
+    except omegaconf.errors.OmegaConfBaseException:
+        return None
+    if key is None:
+        return recipe
+    return recipe[key]
+
+
 class PyTorch(Framework):
     """Handle end-to-end training and deployment of custom PyTorch code."""
 
@@ -551,7 +566,9 @@ def _setup_for_training_recipe(cls, training_recipe, recipe_overrides, kwargs):
         cls.recipe_train_dir = tempfile.TemporaryDirectory(prefix="training_")
         cls.recipe_launcher_dir = tempfile.TemporaryDirectory(prefix="launcher_")
 
-        temp_local_recipe = tempfile.NamedTemporaryFile(prefix="recipe").name
+        temp_local_recipe = tempfile.NamedTemporaryFile(
+            prefix="recipe_original", suffix=".yaml"
+        ).name
         if training_recipe.endswith(".yaml"):
             if os.path.isfile(training_recipe):
                 shutil.copy(training_recipe, temp_local_recipe)
@@ -567,7 +584,7 @@ def _setup_for_training_recipe(cls, training_recipe, recipe_overrides, kwargs):
             _run_clone_command(launcher_repo, cls.recipe_launcher_dir.name)
             recipe = os.path.join(
                 cls.recipe_launcher_dir.name,
-                "recipes-collection",
+                "recipes_collection",
                 "recipes",
                 "training",
                 training_recipe + ".yaml",
@@ -578,6 +595,7 @@ def _setup_for_training_recipe(cls, training_recipe, recipe_overrides, kwargs):
                 raise ValueError(f"Recipe {training_recipe} not found.")
 
         recipe = OmegaConf.load(temp_local_recipe)
+        os.unlink(temp_local_recipe)
 
         if "instance_type" not in kwargs:
             raise ValueError("Must pass instance type to estimator when using training recipes.")
@@ -662,7 +680,26 @@ def _setup_for_training_recipe(cls, training_recipe, recipe_overrides, kwargs):
                 "Ignoring container from training_recipe. Use image_uri arg for estimator."
             )
 
-        OmegaConf.save(config=recipe, f=os.path.join(args["source_dir"], "recipe.yaml"))
+        if not OmegaConf.has_resolver("multiply"):
+            OmegaConf.register_new_resolver("multiply", lambda x, y: x * y, replace=True)
+        if not OmegaConf.has_resolver("divide_ceil"):
+            OmegaConf.register_new_resolver(
+                "divide_ceil", lambda x, y: int(math.ceil(x / y)), replace=True
+            )
+        if not OmegaConf.has_resolver("divide_floor"):
+            OmegaConf.register_new_resolver(
+                "divide_floor", lambda x, y: int(math.floor(x / y)), replace=True
+            )
+        if not OmegaConf.has_resolver("add"):
+            OmegaConf.register_new_resolver("add", lambda *numbers: sum(numbers))
+        final_recipe = _try_resolve_recipe(recipe)
+        if final_recipe is None:
+            final_recipe = _try_resolve_recipe(recipe, "recipes")
+        if final_recipe is None:
+            final_recipe = _try_resolve_recipe(recipe, "training")
+        if final_recipe is None:
+            raise RuntimeError("Could not resolve provided recipe.")
+        OmegaConf.save(config=final_recipe, f=os.path.join(args["source_dir"], "recipe.yaml"))
         args["hyperparameters"] = {"config-path": ".", "config-name": "recipe.yaml"}
 
         return args