sagemaker-hyperpod 3.0.0__py3-none-any.whl

sagemaker/hyperpod/cli/training_utils.py

@@ -0,0 +1,200 @@
+import json
+import pkgutil
+import click
+from typing import Callable, Optional, Mapping, Type
+
+
+def load_schema_for_version(
+    version: str,
+    base_package: str,
+) -> dict:
+    """
+    Load schema.json from the top-level <base_package>.vX_Y_Z package.
+    """
+    ver_pkg = f"{base_package}.v{version.replace('.', '_')}"
+    raw = pkgutil.get_data(ver_pkg, "schema.json")
+    if raw is None:
+        raise click.ClickException(
+            f"Could not load schema.json for version {version} "
+            f"(looked in package {ver_pkg})"
+        )
+    return json.loads(raw)
+
+
+def generate_click_command(
+    *,
+    version_key: Optional[str] = None,
+    schema_pkg: str = "hyperpod_jumpstart_inference_template",
+    registry: Mapping[str, Type] = None,
+) -> Callable:
+    """
+    Decorator factory that:
+      1) Injects click.options from the JSON Schema under `schema_pkg`
+      2) At runtime, pops `version`, builds the flat model from `registry`, calls .to_domain()
+      3) Finally invokes your handler as `func(version, domain_config)`
+    - `version_key`: if given, hard-codes the version (no --version flag injected)
+    - `schema_pkg`: the importable package root to read schema.json from
+    - `registry`: a dict mapping version → flat‐model class, e.g. hyperpod_pytorch_job_template.registry.SCHEMA_REGISTRY
+    """
+    if registry is None:
+        raise ValueError("You must pass a registry mapping version→Model")
+
+    def decorator(func: Callable) -> Callable:
+        # Parser for the single JSON‐dict env var flag
+        def _parse_json_flag(ctx, param, value):
+            if value is None:
+                return None
+            try:
+                return json.loads(value)
+            except json.JSONDecodeError as e:
+                raise click.BadParameter(f"{param.name!r} must be valid JSON: {e}")
+
+        # Parser for list flags
+        def _parse_list_flag(ctx, param, value):
+            if value is None:
+                return None
+            # Remove brackets and split by comma
+            value = value.strip("[]")
+            return [item.strip() for item in value.split(",") if item.strip()]
+
+        # 1) the wrapper click will call
+        def wrapped_func(*args, **kwargs):
+            # extract version
+            version = version_key or kwargs.pop("version", "1.0")
+            debug = kwargs.pop("debug", False)
+
+            # look up the model class
+            Model = registry.get(version)
+            if Model is None:
+                raise click.ClickException(f"Unsupported schema version: {version}")
+
+            # validate & to_domain
+            flat = Model(**kwargs)
+            domain_config = flat.to_domain()
+
+            # call your handler
+            return func(version, debug, domain_config)
+
+        # 2) inject click options from JSON Schema
+        excluded_props = set(["version"])
+        if schema_pkg == "hyperpod_jumpstart_inference_template":
+            wrapped_func = click.option(
+                "--env",
+                callback=_parse_json_flag,
+                type=str,
+                default=None,
+                help=(
+                    "JSON object of environment variables, e.g. "
+                    '\'{"VAR1":"foo","VAR2":"bar"}\''
+                ),
+                metavar="JSON",
+            )(wrapped_func)
+            wrapped_func = click.option(
+                "--resources-limits",
+                callback=_parse_json_flag,
+                help='JSON object of resource limits, e.g. \'{"cpu":"2","memory":"4Gi"}\'',
+                metavar="JSON",
+            )(wrapped_func)
+
+            wrapped_func = click.option(
+                "--resources-requests",
+                callback=_parse_json_flag,
+                help='JSON object of resource requests, e.g. \'{"cpu":"1","memory":"2Gi"}\'',
+                metavar="JSON",
+            )(wrapped_func)
+
+            excluded_props = set(
+                ["version", "env", "resources_limits", "resources_requests"]
+            )
+
+        elif schema_pkg == "hyperpod_pytorch_job_template":
+            wrapped_func = click.option(
+                "--environment",
+                callback=_parse_json_flag,
+                type=str,
+                default=None,
+                help=(
+                    "JSON object of environment variables, e.g. "
+                    '\'{"VAR1":"foo","VAR2":"bar"}\''
+                ),
+                metavar="JSON",
+            )(wrapped_func)
+            wrapped_func = click.option(
+                "--label_selector",
+                callback=_parse_json_flag,
+                help='JSON object of resource limits, e.g. \'{"cpu":"2","memory":"4Gi"}\'',
+                metavar="JSON",
+            )(wrapped_func)
+
+            # Add list options
+            list_params = {
+                "command": "List of command arguments",
+                "args": "List of script arguments, e.g. '[--batch-size, 32, --learning-rate, 0.001]'",
+                "volumes": "List of volumes, e.g. '[vol1, vol2, vol3]'",
+                "persistent_volume_claims": "List of persistent volume claims, e.g. '[pvc1, pvc2]'",
+            }
+
+            for param_name, help_text in list_params.items():
+                wrapped_func = click.option(
+                    f"--{param_name}",
+                    callback=_parse_list_flag,
+                    type=str,
+                    default=None,
+                    help=help_text,
+                    metavar="LIST",
+                )(wrapped_func)
+
+            excluded_props = set(
+                [
+                    "version",
+                    "environment",
+                    "label_selector",
+                    "command",
+                    "args",
+                    "volumes",
+                    "persistent_volume_claims",
+                ]
+            )
+
+        schema = load_schema_for_version(version_key or "1.0", schema_pkg)
+        props = schema.get("properties", {})
+        reqs = set(schema.get("required", []))
+
+        # reverse so flags appear in the same order as in schema.json
+        for name, spec in reversed(list(props.items())):
+            if name in excluded_props:
+                continue
+
+            # type inference
+            if "enum" in spec:
+                ctype = click.Choice(spec["enum"])
+            elif spec.get("type") == "integer":
+                ctype = int
+            elif spec.get("type") == "number":
+                ctype = float
+            elif spec.get("type") == "boolean":
+                ctype = bool
+            else:
+                ctype = str
+
+            wrapped_func = click.option(
+                f"--{name.replace('_','-')}",
+                required=(name in reqs),
+                default=spec.get("default", None),
+                show_default=("default" in spec),
+                type=ctype,
+                help=spec.get("description", ""),
+            )(wrapped_func)
+
+        # 3) if no hard-coded version_key, inject the top-level --version flag
+        if version_key is None:
+            wrapped_func = click.option(
+                "--version",
+                default="1.0",
+                show_default=True,
+                help="Schema version to use",
+            )(wrapped_func)
+
+        return wrapped_func
+
+    return decorator

sagemaker/hyperpod/cli/utils.py

@@ -0,0 +1,183 @@
+# 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.
+import logging
+import re
+import json
+
+import boto3
+import botocore
+from botocore.config import Config
+
+from sagemaker.hyperpod.cli.constants.command_constants import (
+    GENERATED_LAUNCHER_CONFIG_FILE_PATH,
+    HYPERPOD_CLUSTER_CONTEXT_FILE_NAME,
+)
+
+
+def get_name_from_arn(arn: str) -> str:
+    """
+    Parse the EKS cluster name from an EKS ARN.
+
+    Args:
+        arn (str): The ARN of the EKS cluster.
+
+    Returns: str: The name of the EKS cluster if parsing is
+    successful, otherwise raise RuntimeError.
+    """
+    # Define the regex pattern to match the EKS ARN and capture the cluster name
+    pattern = r"arn:aws:eks:[\w-]+:\d+:cluster/([\w-]+)"
+    match = re.match(pattern, arn)
+
+    if match:
+        return match.group(1)
+    else:
+        raise RuntimeError("cannot get EKS cluster name")
+
+
+def setup_logger(
+    name: str,
+    logging_level: int = logging.ERROR,
+) -> logging.Logger:
+    """
+    Set up a logger with a console handler and a formatter.
+
+    Args:
+        name (str): The name of the logger.
+        logging_level (int): The logging level.
+
+    Returns:
+        logging.Logger: The configured logger instance.
+    """
+    # Create a logger
+    logger = logging.getLogger(name)
+    logger.setLevel(logging_level)  # Set the log level to ERROR
+
+    # Create a console handler
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging_level)  # Set the log level for the handler
+
+    # Create a formatter and set it for the handler
+    formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+    console_handler.setFormatter(formatter)
+
+    # Add the handler to the logger
+    logger.addHandler(console_handler)
+
+    return logger
+
+
+def set_logging_level(
+    logger: logging.Logger,
+    logging_level: int,
+):
+    logger.setLevel(logging_level)
+    logger.handlers[0].setLevel(logging_level)
+
+
+def get_sagemaker_client(
+    session: boto3.Session, config: Config = None
+) -> botocore.client.BaseClient:
+    return session.client(
+        service_name="sagemaker",
+        config=config,
+    )
+
+
+def store_current_hyperpod_context(data):
+    with open(
+        GENERATED_LAUNCHER_CONFIG_FILE_PATH + HYPERPOD_CLUSTER_CONTEXT_FILE_NAME,
+        "w",
+    ) as hyperpod_current_context:
+        hyperpod_current_context.write(json.dumps(data, indent=4, default=str))
+
+
+def _retrieve_current_hyperpod_context():
+    with open(
+        GENERATED_LAUNCHER_CONFIG_FILE_PATH + HYPERPOD_CLUSTER_CONTEXT_FILE_NAME,
+        "r",
+    ) as file:
+        return json.load(file)
+
+
+def _validate_link(console_url):
+    pattern = "https:\/\/([a-z0-9-]+).console.aws.amazon.com\/sagemaker\/home\?region=([a-z0-9-]+)#\/cluster-management\/([a-zA-Z0-9-]+)"
+    match = re.match(pattern, console_url)
+    if match:
+        return True
+    else:
+        return False
+
+
+def validate_region_and_cluster_name(region, cluster_name):
+    output = False
+    region_char_list = region.split("-")
+
+    if len(region_char_list) != 3:
+        return False
+
+    region_prefix_match = re.match("[a-z]+", region_char_list[0])
+    region_match = re.match("[a-z]+", region_char_list[1])
+    region_suffix_match = re.match("[0-9]+", region_char_list[2])
+
+    region_prefix_length = len(region_char_list[0])
+    region_length = len(region_char_list[1])
+    region_suffix_length = len(region_char_list[2])
+
+    cluster_name_match = re.match("[a-zA-Z0-9-]+", cluster_name)
+
+    if (
+        region_prefix_match
+        and region_match
+        and region_suffix_match
+        and region_prefix_length == 2
+        and region_suffix_length == 1
+        and region_length >= 4
+        and region_length < 10
+        and cluster_name_match
+        and len(cluster_name) >= 1
+        and len(cluster_name) <= 63
+    ):
+        output = True
+    return output
+
+
+def get_cluster_console_url():
+    hyperpod_context_cluster = _retrieve_current_hyperpod_context()
+    console_url = None
+    if (
+        hyperpod_context_cluster
+        and hyperpod_context_cluster.get("ClusterArn")
+        and hyperpod_context_cluster.get("ClusterName")
+    ):
+        region = hyperpod_context_cluster.get("ClusterArn").split(":")[3]
+        cluster_name = hyperpod_context_cluster.get("ClusterName")
+
+        console_url = (
+            f"https://{region}.console.aws.amazon.com/sagemaker/"
+            f"home?region={region}#/cluster-management/{cluster_name}"
+        )
+        if _validate_link(console_url) and validate_region_and_cluster_name(region, cluster_name):
+            return console_url
+    return None
+
+def get_eks_cluster_name():
+    hyperpod_context_cluster = _retrieve_current_hyperpod_context()
+    eks_cluster_arn = hyperpod_context_cluster.get("Orchestrator", {}).get("Eks", {}).get("ClusterArn", '')
+    return eks_cluster_arn.split('cluster/')[-1]
+
+def get_hyperpod_cluster_region():
+    hyperpod_context_cluster = _retrieve_current_hyperpod_context()
+    return hyperpod_context_cluster.get("ClusterArn").split(":")[3]

sagemaker/hyperpod/cli/validators/__init__.py

@@ -0,0 +1,12 @@
+# 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.

sagemaker/hyperpod/cli/validators/cluster_validator.py

@@ -0,0 +1,52 @@
+# 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.
+from typing import Optional
+
+import boto3
+from botocore.exceptions import ClientError
+
+from sagemaker.hyperpod.cli.utils import setup_logger
+from sagemaker.hyperpod.cli.validators.validator import (
+    Validator,
+)
+
+logger = setup_logger(__name__)
+
+
+class ClusterValidator(Validator):
+    def __init__(self):
+        super().__init__()
+
+    def validate_cluster_and_get_eks_arn(
+        self, cluster_name: str, sm_client: boto3.client
+    ) -> Optional[str]:
+        try:
+            hp_cluster_details = sm_client.describe_cluster(ClusterName=cluster_name)
+            if (
+                "Orchestrator" not in hp_cluster_details
+                or "Eks" not in hp_cluster_details["Orchestrator"]
+            ):
+                logger.warning(
+                    f"HyperPod cluster {cluster_name} Orchestrator not exist or is not Eks."
+                )
+                return None
+            return hp_cluster_details["Orchestrator"]["Eks"]["ClusterArn"]
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "ResourceNotFoundException":
+                logger.error(f"HyperPod cluster {cluster_name} not found.")
+            else:
+                logger.error(f"Validate HyperPod cluster {cluster_name} error: {e}")
+            return None
+        except Exception as e:
+            logger.error(f"Validate HyperPod cluster {cluster_name} error: {e}")
+            return None