sagemaker-core 1.0.47__py3-none-any.whl → 2.1.1__py3-none-any.whl

sagemaker/core/modules/train/container_drivers/scripts/environment.py

@@ -0,0 +1,305 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""This module is used to define the environment variables for the training job container."""
+from __future__ import absolute_import
+
+from typing import Dict, Any
+import multiprocessing
+import subprocess
+import json
+import os
+import sys
+from pathlib import Path
+import logging
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from common.utils import (  # noqa: E402 # pylint: disable=C0413,E0611
+    safe_serialize,
+    safe_deserialize,
+    read_distributed_json,
+    read_source_code_json,
+)
+
+# Initialize logger
+SM_LOG_LEVEL = os.environ.get("SM_LOG_LEVEL", 20)
+logger = logging.getLogger(__name__)
+console_handler = logging.StreamHandler(sys.stdout)
+logger.addHandler(console_handler)
+logger.setLevel(int(SM_LOG_LEVEL))
+
+SM_MODEL_DIR = "/opt/ml/model"
+
+SM_INPUT_DIR = "/opt/ml/input"
+SM_INPUT_DATA_DIR = "/opt/ml/input/data"
+SM_INPUT_CONFIG_DIR = "/opt/ml/input/config"
+
+SM_OUTPUT_DIR = "/opt/ml/output"
+SM_OUTPUT_FAILURE = "/opt/ml/output/failure"
+SM_OUTPUT_DATA_DIR = "/opt/ml/output/data"
+SM_SOURCE_DIR_PATH = "/opt/ml/input/data/code"
+SM_DISTRIBUTED_DRIVER_DIR_PATH = "/opt/ml/input/data/sm_drivers/distributed_drivers"
+
+SM_MASTER_ADDR = "algo-1"
+SM_MASTER_PORT = 7777
+
+RESOURCE_CONFIG = f"{SM_INPUT_CONFIG_DIR}/resourceconfig.json"
+INPUT_DATA_CONFIG = f"{SM_INPUT_CONFIG_DIR}/inputdataconfig.json"
+HYPERPARAMETERS_CONFIG = f"{SM_INPUT_CONFIG_DIR}/hyperparameters.json"
+
+ENV_OUTPUT_FILE = "/opt/ml/input/sm_training.env"
+
+SENSITIVE_KEYWORDS = ["SECRET", "PASSWORD", "KEY", "TOKEN", "PRIVATE", "CREDS", "CREDENTIALS"]
+HIDDEN_VALUE = "******"
+
+
+def num_cpus() -> int:
+    """Return the number of CPUs available in the current container.
+
+    Returns:
+        int: Number of CPUs available in the current container.
+    """
+    return multiprocessing.cpu_count()
+
+
+def num_gpus() -> int:
+    """Return the number of GPUs available in the current container.
+
+    Returns:
+        int: Number of GPUs available in the current container.
+    """
+    try:
+        cmd = ["nvidia-smi", "--list-gpus"]
+        output = subprocess.check_output(cmd).decode("utf-8")
+        return sum(1 for line in output.splitlines() if line.startswith("GPU "))
+    except (OSError, subprocess.CalledProcessError):
+        logger.info("No GPUs detected (normal if no gpus installed)")
+        return 0
+
+
+def num_neurons() -> int:
+    """Return the number of neuron cores available in the current container.
+
+    Returns:
+        int: Number of Neuron Cores available in the current container.
+    """
+    try:
+        cmd = ["neuron-ls", "-j"]
+        output = subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8")
+        j = json.loads(output)
+        neuron_cores = 0
+        for item in j:
+            neuron_cores += item.get("nc_count", 0)
+        logger.info("Found %s neurons on this instance", neuron_cores)
+        return neuron_cores
+    except OSError:
+        logger.info("No Neurons detected (normal if no neurons installed)")
+        return 0
+    except subprocess.CalledProcessError as e:
+        if e.output is not None:
+            try:
+                msg = e.output.decode("utf-8").partition("error=")[2]
+                logger.info(
+                    "No Neurons detected (normal if no neurons installed). \
+                    If neuron installed then %s",
+                    msg,
+                )
+            except AttributeError:
+                logger.info("No Neurons detected (normal if no neurons installed)")
+        else:
+            logger.info("No Neurons detected (normal if no neurons installed)")
+
+        return 0
+
+
+def deserialize_hyperparameters(hyperparameters: Dict[str, str]) -> Dict[str, Any]:
+    """Deserialize hyperparameters from string to their original types.
+
+    Args:
+        hyperparameters (Dict[str, str]): Hyperparameters as strings.
+
+    Returns:
+        Dict[str, Any]: Hyperparameters as their original types.
+    """
+    deserialized_hyperparameters = {}
+    for key, value in hyperparameters.items():
+        deserialized_hyperparameters[key] = safe_deserialize(value)
+    return deserialized_hyperparameters
+
+
+def set_env(
+    resource_config: Dict[str, Any],
+    input_data_config: Dict[str, Any],
+    hyperparameters_config: Dict[str, Any],
+    output_file: str = ENV_OUTPUT_FILE,
+):
+    """Set environment variables for the training job container.
+
+    Args:
+        resource_config (Dict[str, Any]): Resource configuration for the training job.
+        input_data_config (Dict[str, Any]): Input data configuration for the training job.
+        hyperparameters_config (Dict[str, Any]): Hyperparameters configuration for the training job.
+        output_file (str): Output file to write the environment variables.
+    """
+    # Constants
+    env_vars = {
+        "SM_MODEL_DIR": SM_MODEL_DIR,
+        "SM_INPUT_DIR": SM_INPUT_DIR,
+        "SM_INPUT_DATA_DIR": SM_INPUT_DATA_DIR,
+        "SM_INPUT_CONFIG_DIR": SM_INPUT_CONFIG_DIR,
+        "SM_OUTPUT_DIR": SM_OUTPUT_DIR,
+        "SM_OUTPUT_FAILURE": SM_OUTPUT_FAILURE,
+        "SM_OUTPUT_DATA_DIR": SM_OUTPUT_DATA_DIR,
+        "SM_LOG_LEVEL": SM_LOG_LEVEL,
+        "SM_MASTER_ADDR": SM_MASTER_ADDR,
+        "SM_MASTER_PORT": SM_MASTER_PORT,
+    }
+
+    # SourceCode and DistributedConfig Environment Variables
+    source_code = read_source_code_json()
+    if source_code:
+        env_vars["SM_SOURCE_DIR"] = SM_SOURCE_DIR_PATH
+        env_vars["SM_ENTRY_SCRIPT"] = source_code.get("entry_script", "")
+
+    distributed = read_distributed_json()
+    if distributed:
+        env_vars["SM_DISTRIBUTED_DRIVER_DIR"] = SM_DISTRIBUTED_DRIVER_DIR_PATH
+        env_vars["SM_DISTRIBUTED_CONFIG"] = distributed
+
+    # Data Channels
+    channels = list(input_data_config.keys())
+    for channel in channels:
+        env_vars[f"SM_CHANNEL_{channel.upper()}"] = f"{SM_INPUT_DATA_DIR}/{channel}"
+    env_vars["SM_CHANNELS"] = channels
+
+    # Hyperparameters
+    hps = deserialize_hyperparameters(hyperparameters_config)
+    for key, value in hps.items():
+        key_upper = key.replace("-", "_").upper()
+        env_vars[f"SM_HP_{key_upper}"] = value
+    env_vars["SM_HPS"] = hps
+
+    # Host Variables
+    current_host = resource_config["current_host"]
+    current_instance_type = resource_config["current_instance_type"]
+    hosts = resource_config["hosts"]
+    sorted_hosts = sorted(hosts)
+
+    env_vars["SM_CURRENT_HOST"] = current_host
+    env_vars["SM_CURRENT_INSTANCE_TYPE"] = current_instance_type
+    env_vars["SM_HOSTS"] = sorted_hosts
+    env_vars["SM_NETWORK_INTERFACE_NAME"] = resource_config["network_interface_name"]
+    env_vars["SM_HOST_COUNT"] = len(sorted_hosts)
+    env_vars["SM_CURRENT_HOST_RANK"] = sorted_hosts.index(current_host)
+
+    env_vars["SM_NUM_CPUS"] = num_cpus()
+    env_vars["SM_NUM_GPUS"] = num_gpus()
+    env_vars["SM_NUM_NEURONS"] = num_neurons()
+
+    # Misc.
+    env_vars["SM_RESOURCE_CONFIG"] = resource_config
+    env_vars["SM_INPUT_DATA_CONFIG"] = input_data_config
+
+    # All Training Environment Variables
+    env_vars["SM_TRAINING_ENV"] = {
+        "channel_input_dirs": {
+            channel: env_vars[f"SM_CHANNEL_{channel.upper()}"] for channel in channels
+        },
+        "current_host": env_vars["SM_CURRENT_HOST"],
+        "current_instance_type": env_vars["SM_CURRENT_INSTANCE_TYPE"],
+        "hosts": env_vars["SM_HOSTS"],
+        "master_addr": env_vars["SM_MASTER_ADDR"],
+        "master_port": env_vars["SM_MASTER_PORT"],
+        "hyperparameters": env_vars["SM_HPS"],
+        "input_data_config": input_data_config,
+        "input_config_dir": env_vars["SM_INPUT_CONFIG_DIR"],
+        "input_data_dir": env_vars["SM_INPUT_DATA_DIR"],
+        "input_dir": env_vars["SM_INPUT_DIR"],
+        "job_name": os.environ["TRAINING_JOB_NAME"],
+        "log_level": env_vars["SM_LOG_LEVEL"],
+        "model_dir": env_vars["SM_MODEL_DIR"],
+        "network_interface_name": env_vars["SM_NETWORK_INTERFACE_NAME"],
+        "num_cpus": env_vars["SM_NUM_CPUS"],
+        "num_gpus": env_vars["SM_NUM_GPUS"],
+        "num_neurons": env_vars["SM_NUM_NEURONS"],
+        "output_data_dir": env_vars["SM_OUTPUT_DATA_DIR"],
+        "resource_config": env_vars["SM_RESOURCE_CONFIG"],
+    }
+    with open(output_file, "w") as f:
+        for key, value in env_vars.items():
+            f.write(f"export {key}='{safe_serialize(value)}'\n")
+
+    logger.info("Environment Variables:")
+    log_env_variables(env_vars_dict=env_vars)
+
+
+def mask_sensitive_info(data):
+    """Recursively mask sensitive information in a dictionary."""
+    if isinstance(data, dict):
+        for k, v in data.items():
+            if isinstance(v, dict):
+                data[k] = mask_sensitive_info(v)
+            elif isinstance(v, str) and any(
+                keyword.lower() in k.lower() for keyword in SENSITIVE_KEYWORDS
+            ):
+                data[k] = HIDDEN_VALUE
+    return data
+
+
+def log_key_value(key: str, value: str):
+    """Log a key-value pair, masking sensitive values if necessary."""
+    if any(keyword.lower() in key.lower() for keyword in SENSITIVE_KEYWORDS):
+        logger.info("%s=%s", key, HIDDEN_VALUE)
+    elif isinstance(value, dict):
+        masked_value = mask_sensitive_info(value)
+        logger.info("%s=%s", key, json.dumps(masked_value))
+    else:
+        try:
+            decoded_value = json.loads(value)
+            if isinstance(decoded_value, dict):
+                masked_value = mask_sensitive_info(decoded_value)
+                logger.info("%s=%s", key, json.dumps(masked_value))
+            else:
+                logger.info("%s=%s", key, decoded_value)
+        except (json.JSONDecodeError, TypeError):
+            logger.info("%s=%s", key, value)
+
+
+def log_env_variables(env_vars_dict: Dict[str, Any]):
+    """Log Environment Variables from the environment and an env_vars_dict."""
+    for key, value in os.environ.items():
+        log_key_value(key, value)
+
+    for key, value in env_vars_dict.items():
+        log_key_value(key, value)
+
+
+def main():
+    """Main function to set the environment variables for the training job container."""
+    with open(RESOURCE_CONFIG, "r") as f:
+        resource_config = json.load(f)
+    with open(INPUT_DATA_CONFIG, "r") as f:
+        input_data_config = json.load(f)
+    with open(HYPERPARAMETERS_CONFIG, "r") as f:
+        hyperparameters_config = json.load(f)
+
+    set_env(
+        resource_config=resource_config,
+        input_data_config=input_data_config,
+        hyperparameters_config=hyperparameters_config,
+        output_file=ENV_OUTPUT_FILE,
+    )
+
+
+if __name__ == "__main__":
+    main()

sagemaker/core/modules/train/sm_recipes/utils.py

@@ -0,0 +1,330 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Utility functions for SageMaker training recipes."""
+from __future__ import absolute_import
+
+import math
+import os
+import json
+import shutil
+import tempfile
+from urllib.request import urlretrieve
+from typing import Dict, Any, Optional, Tuple
+
+import omegaconf
+from omegaconf import OmegaConf, dictconfig
+
+from sagemaker.core.image_uris import retrieve
+
+from sagemaker.core.modules import logger
+from sagemaker.core.modules.utils import _run_clone_command_silent
+from sagemaker.core.modules.configs import Compute, SourceCode
+from sagemaker.core.modules.distributed import Torchrun, SMP
+
+
+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]
+
+
+def _determine_device_type(instance_type: str) -> str:
+    """Determine device type (gpu, cpu, trainium) based on instance type."""
+    instance_family = instance_type.split(".")[1]
+    if instance_family.startswith(("p", "g")):
+        return "gpu"
+    if instance_family.startswith("trn"):
+        return "trainium"
+    return "cpu"
+
+
+def _load_recipes_cfg() -> str:
+    """Load training recipes configuration json."""
+    training_recipes_cfg_filename = os.path.join(os.path.dirname(__file__), "training_recipes.json")
+    with open(training_recipes_cfg_filename) as training_recipes_cfg_file:
+        training_recipes_cfg = json.load(training_recipes_cfg_file)
+    return training_recipes_cfg
+
+
+def _load_base_recipe(
+    training_recipe: str,
+    recipe_overrides: Optional[Dict[str, Any]] = None,
+    training_recipes_cfg: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    """Load recipe and apply overrides."""
+    if recipe_overrides is None:
+        recipe_overrides = dict()
+
+    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)
+        else:
+            try:
+                urlretrieve(training_recipe, temp_local_recipe)
+            except Exception as e:
+                raise ValueError(
+                    f"Could not fetch the provided recipe {training_recipe}: exception {str(e)}"
+                )
+    else:
+        recipe_launcher_dir = tempfile.TemporaryDirectory(prefix="launcher_")
+
+        launcher_repo = os.environ.get("TRAINING_LAUNCHER_GIT", None) or training_recipes_cfg.get(
+            "launcher_repo"
+        )
+        _run_clone_command_silent(launcher_repo, recipe_launcher_dir.name)
+
+        recipe = os.path.join(
+            recipe_launcher_dir.name,
+            "recipes_collection",
+            "recipes",
+            training_recipe + ".yaml",
+        )
+        if os.path.isfile(recipe):
+            shutil.copy(recipe, temp_local_recipe)
+        else:
+            raise ValueError(f"Recipe {training_recipe} not found.")
+
+    recipe = OmegaConf.load(temp_local_recipe)
+    os.unlink(temp_local_recipe)
+    recipe = OmegaConf.merge(recipe, recipe_overrides)
+    return recipe
+
+
+def _register_custom_resolvers():
+    """Register custom resolvers for OmegaConf."""
+    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))
+
+
+def _get_trainining_recipe_gpu_model_name_and_script(model_type: str):
+    """Get the model base name and script for the training recipe."""
+
+    model_type_to_script = {
+        "llama_v3": ("llama", "llama_pretrain.py"),
+        "mistral": ("mistral", "mistral_pretrain.py"),
+        "mixtral": ("mixtral", "mixtral_pretrain.py"),
+        "deepseek": ("deepseek", "deepseek_pretrain.py"),
+    }
+
+    for key in model_type_to_script:
+        if model_type.startswith(key):
+            model_type = key
+            break
+
+    if model_type not in model_type_to_script:
+        raise ValueError(f"Model type {model_type} not supported")
+
+    return model_type_to_script[model_type][0], model_type_to_script[model_type][1]
+
+
+def _configure_gpu_args(
+    training_recipes_cfg: Dict[str, Any],
+    region_name: str,
+    recipe: OmegaConf,
+    recipe_train_dir: tempfile.TemporaryDirectory,
+) -> Dict[str, Any]:
+    """Configure arguments specific to GPU."""
+    source_code = SourceCode()
+    args = dict()
+
+    adapter_repo = os.environ.get("TRAINING_ADAPTER_GIT", None) or training_recipes_cfg.get(
+        "adapter_repo"
+    )
+    _run_clone_command_silent(adapter_repo, recipe_train_dir.name)
+
+    if "model" not in recipe:
+        raise ValueError("Supplied recipe does not contain required field model.")
+    if "model_type" not in recipe["model"]:
+        raise ValueError("Supplied recipe does not contain required field model_type.")
+    model_type = recipe["model"]["model_type"]
+
+    model_base_name, script = _get_trainining_recipe_gpu_model_name_and_script(model_type)
+
+    source_code.source_dir = os.path.join(recipe_train_dir.name, "examples", model_base_name)
+    source_code.entry_script = script
+
+    gpu_image_cfg = training_recipes_cfg.get("gpu_image")
+    if isinstance(gpu_image_cfg, str):
+        training_image = gpu_image_cfg
+    else:
+        training_image = retrieve(
+            gpu_image_cfg.get("framework"),
+            region=region_name,
+            version=gpu_image_cfg.get("version"),
+            image_scope="training",
+            **gpu_image_cfg.get("additional_args"),
+        )
+
+    # Setting dummy parameters for now
+    torch_distributed = Torchrun(smp=SMP(random_seed="123456"))
+    args.update(
+        {
+            "source_code": source_code,
+            "training_image": training_image,
+            "distributed": torch_distributed,
+        }
+    )
+    return args
+
+
+def _configure_trainium_args(
+    training_recipes_cfg: Dict[str, Any],
+    region_name: str,
+    recipe_train_dir: tempfile.TemporaryDirectory,
+) -> Dict[str, Any]:
+    """Configure arguments specific to Trainium."""
+    source_code = SourceCode()
+    args = dict()
+
+    _run_clone_command_silent(training_recipes_cfg.get("neuron_dist_repo"), recipe_train_dir.name)
+
+    source_code.source_dir = os.path.join(recipe_train_dir.name, "examples")
+    source_code.entry_script = "training_orchestrator.py"
+    neuron_image_cfg = training_recipes_cfg.get("neuron_image")
+    if isinstance(neuron_image_cfg, str):
+        training_image = neuron_image_cfg
+    else:
+        training_image = retrieve(
+            neuron_image_cfg.get("framework"),
+            region=region_name,
+            version=neuron_image_cfg.get("version"),
+            image_scope="training",
+            **neuron_image_cfg.get("additional_args"),
+        )
+
+    args.update(
+        {
+            "source_code": source_code,
+            "training_image": training_image,
+            "distributed": Torchrun(),
+        }
+    )
+    return args
+
+
+def _get_args_from_recipe(
+    training_recipe: str,
+    compute: Compute,
+    region_name: str,
+    recipe_overrides: Optional[Dict[str, Any]],
+    requirements: Optional[str],
+) -> Tuple[Dict[str, Any], tempfile.TemporaryDirectory]:
+    """Get arguments for ModelTrainer from a training recipe.
+
+    Returns a dictionary of arguments to be used with ModelTrainer like:
+    ```python
+    {
+        "source_code": SourceCode,
+        "training_image": str,
+        "distributed": DistributedConfig,
+        "compute": Compute,
+        "hyperparameters": Dict[str, Any],
+    }
+    ```
+
+    Args:
+        training_recipe (str):
+            Name of the training recipe or path to the recipe file.
+        compute (Compute):
+            Compute configuration for training.
+        region_name (str):
+            Name of the AWS region.
+        recipe_overrides (Optional[Dict[str, Any]]):
+            Overrides for the training recipe.
+        requirements (Optional[str]):
+            Path to the requirements file.
+    """
+    if compute.instance_type is None:
+        raise ValueError("Must set `instance_type` in compute when using training recipes.")
+
+    training_recipes_cfg = _load_recipes_cfg()
+    recipe = _load_base_recipe(training_recipe, recipe_overrides, training_recipes_cfg)
+
+    if "trainer" not in recipe:
+        raise ValueError("Supplied recipe does not contain required field trainer.")
+
+    # Set instance_count
+    if compute.instance_count and "num_nodes" in recipe["trainer"]:
+        logger.warning(
+            f"Using Compute to set instance_count:\n{compute}."
+            "\nIgnoring trainer -> num_nodes in recipe."
+        )
+    if compute.instance_count is None:
+        if "num_nodes" not in recipe["trainer"]:
+            raise ValueError(
+                "Must provide Compute with instance_count or" " set trainer -> num_nodes in recipe."
+            )
+        compute.instance_count = recipe["trainer"]["num_nodes"]
+
+    if requirements and not os.path.isfile(requirements):
+        raise ValueError(f"Recipe requirements file {requirements} not found.")
+
+    # Get Training Image, SourceCode, and distributed args
+    device_type = _determine_device_type(compute.instance_type)
+    recipe_train_dir = tempfile.TemporaryDirectory(prefix="training_")
+    if device_type == "gpu":
+        args = _configure_gpu_args(training_recipes_cfg, region_name, recipe, recipe_train_dir)
+    elif device_type == "trainium":
+        args = _configure_trainium_args(training_recipes_cfg, region_name, recipe_train_dir)
+    else:
+        raise ValueError(f"Devices of type {device_type} are not supported with training recipes.")
+
+    _register_custom_resolvers()
+
+    # Resolve Final Recipe
+    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.")
+
+    # Save Final Recipe to source_dir
+    OmegaConf.save(
+        config=final_recipe, f=os.path.join(args["source_code"].source_dir, "recipe.yaml")
+    )
+
+    # If recipe_requirements is provided, copy it to source_dir
+    if requirements:
+        shutil.copy(requirements, args["source_code"].source_dir)
+        args["source_code"].requirements = os.path.basename(requirements)
+
+    # Update args with compute and hyperparameters
+    args.update(
+        {
+            "compute": compute,
+            "hyperparameters": {"config-path": ".", "config-name": "recipe.yaml"},
+        }
+    )
+
+    return args, recipe_train_dir

sagemaker/core/modules/types.py

@@ -0,0 +1,19 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Types module."""
+from __future__ import absolute_import
+
+from typing import Union
+from sagemaker.core.modules.configs import S3DataSource, FileSystemDataSource
+
+DataSourceType = Union[str, S3DataSource, FileSystemDataSource]