sagemaker-hyperpod 3.0.0__py3-none-any.whl

sagemaker/hyperpod/cli/hyp_cli.py

@@ -0,0 +1,132 @@
+import click
+import yaml
+import json
+import os
+import subprocess
+from pydantic import BaseModel, ValidationError, Field
+from typing import Optional
+
+from sagemaker.hyperpod.cli.commands.cluster import list_cluster, set_cluster_context, get_cluster_context, \
+    get_monitoring
+from sagemaker.hyperpod.cli.commands.training import (
+    pytorch_create,
+    list_jobs,
+    pytorch_describe,
+    pytorch_delete,
+    pytorch_list_pods,
+    pytorch_get_logs,
+)
+from sagemaker.hyperpod.cli.commands.inference import (
+    js_create,
+    custom_create,
+    custom_invoke,
+    js_list,
+    custom_list,
+    js_describe,
+    custom_describe,
+    js_delete,
+    custom_delete,
+    js_list_pods,
+    custom_list_pods,
+    js_get_logs,
+    custom_get_logs,
+    js_get_operator_logs,
+    custom_get_operator_logs,
+)
+
+
+@click.group()
+def cli():
+    pass
+
+
+class CLICommand(click.Group):
+    pass
+
+
+@cli.group(cls=CLICommand)
+def create():
+    """Create a jumpstart model endpoint, a custom model endpoint, or a pytorch job."""
+    pass
+
+
+@cli.group(cls=CLICommand)
+def list():
+    """List all jumpstart model endpoints, custom model endpoints, or pytorch jobs."""
+    pass
+
+
+@cli.group(cls=CLICommand)
+def describe():
+    """Describe a jumpstart model endpoint, a custom model endpoint, or a pytorch job."""
+    pass
+
+
+@cli.group(cls=CLICommand)
+def delete():
+    """Delete a jumpstart model endpoint, a custom model endpoint, or a pytorch job."""
+    pass
+
+
+@cli.group(cls=CLICommand)
+def list_pods():
+    """List all pods for jumpstart model endpoint, custom model endpoint or pytorch jobs."""
+    pass
+
+
+@cli.group(cls=CLICommand)
+def get_logs():
+    """Get specific pod logs for a jumpstart model endpoint, custom model endpoint or pytorch job."""
+    pass
+
+
+@cli.group(cls=CLICommand)
+def invoke():
+    """Invoke a jumpstart model endpoint or a custom model endpoint."""
+    pass
+
+
+@cli.group(cls=CLICommand)
+def get_operator_logs():
+    """Get operator logs for jumpstart model endpoint, or custom model endpoint."""
+    pass
+
+
+create.add_command(pytorch_create)
+create.add_command(js_create)
+create.add_command(custom_create)
+
+list.add_command(list_jobs)
+list.add_command(js_list)
+list.add_command(custom_list)
+
+describe.add_command(pytorch_describe)
+describe.add_command(js_describe)
+describe.add_command(custom_describe)
+
+delete.add_command(pytorch_delete)
+delete.add_command(js_delete)
+delete.add_command(custom_delete)
+
+list_pods.add_command(pytorch_list_pods)
+list_pods.add_command(js_list_pods)
+list_pods.add_command(custom_list_pods)
+
+get_logs.add_command(pytorch_get_logs)
+get_logs.add_command(js_get_logs)
+get_logs.add_command(custom_get_logs)
+
+get_operator_logs.add_command(js_get_operator_logs)
+get_operator_logs.add_command(custom_get_operator_logs)
+
+invoke.add_command(custom_invoke)
+invoke.add_command(custom_invoke, name="hyp-jumpstart-endpoint")
+
+cli.add_command(list_cluster)
+cli.add_command(set_cluster_context)
+cli.add_command(get_cluster_context)
+cli.add_command(get_monitoring)
+
+
+if __name__ == "__main__":
+    cli()

sagemaker/hyperpod/cli/inference_utils.py

@@ -0,0 +1,121 @@
+import json
+import pkgutil
+import click
+from typing import Callable, Optional, Mapping, Type
+
+
+def load_schema_for_version(version: str, schema_pkg: str) -> dict:
+    ver_pkg = f"{schema_pkg}.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}")
+    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:
+    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}")
+
+        # 1) the wrapper click actually invokes
+        def wrapped_func(*args, **kwargs):
+            namespace = kwargs.pop("namespace", None)
+            version = version_key or kwargs.pop("version", "1.0")
+
+            Model = registry.get(version)
+            if Model is None:
+                raise click.ClickException(f"Unsupported schema version: {version}")
+
+            flat = Model(**kwargs)
+            domain = flat.to_domain()
+            return func(namespace, version, domain)
+
+        # 2) inject the special JSON‐env flag before everything else
+        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(
+            "--dimensions",
+            callback=_parse_json_flag,
+            type=str,
+            default=None,
+            help=("JSON object of dimensions, 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)
+
+        # 3) auto-inject all schema.json fields
+        schema = load_schema_for_version(version_key or "1.0", schema_pkg)
+        props = schema.get("properties", {})
+        reqs = set(schema.get("required", []))
+
+        for name, spec in reversed(list(props.items())):
+            if name in (
+                "version",
+                "env",
+                "dimensions",
+                "resources_limits",
+                "resources_requests",
+            ):
+                continue
+
+            # infer click type
+            if "enum" in spec:
+                ctype = click.Choice(spec["enum"])
+            if 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)
+
+        return wrapped_func
+
+    return decorator

sagemaker/hyperpod/cli/service/__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/service/cancel_training_job.py

@@ -0,0 +1,70 @@
+# 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 subprocess
+
+from sagemaker.hyperpod.cli.clients.kubernetes_client import (
+    KubernetesClient,
+)
+from kubernetes.client.rest import ApiException
+
+from sagemaker.hyperpod.cli.constants.pytorch_constants import (
+    PYTORCH_CUSTOM_OBJECT_GROUP,
+    PYTORCH_CUSTOM_OBJECT_PLURAL
+)
+from sagemaker.hyperpod.cli.service.discover_namespaces import DiscoverNamespaces
+from kubernetes.client import (
+    V1ResourceAttributes
+)
+
+class CancelTrainingJob:
+    def __init__(self):
+        return
+
+    def cancel_training_job(self, job_name: str, namespace: Optional[str]):
+        """
+        Cancel training job provided by the user in the specified namespace.
+        If namespace is not provided job is canceled from the default namespace in user context
+        """
+
+        k8s_client = KubernetesClient()
+
+        if not namespace:
+            resource_attributes_template = V1ResourceAttributes(
+                verb="delete",
+                group=PYTORCH_CUSTOM_OBJECT_GROUP,
+                resource=PYTORCH_CUSTOM_OBJECT_PLURAL,
+            )
+            namespace = DiscoverNamespaces().discover_accessible_namespace(
+                resource_attributes_template
+            )
+        try:
+            result = k8s_client.delete_training_job(
+                job_name=job_name, namespace=namespace
+            )
+        except ApiException as e:
+            raise RuntimeError(f"Unexpected API error: {e.reason} ({e.status})")
+
+        helm_chart_cleanup_command = [
+            "helm",
+            "uninstall",
+            job_name,
+            "--namespace",
+            namespace,
+        ]
+
+        if result.get("status") and result.get("status") == "Success":
+            subprocess.run(helm_chart_cleanup_command, capture_output=True, text=True)
+            return None
+        else:
+            return result

sagemaker/hyperpod/cli/service/discover_namespaces.py

@@ -0,0 +1,108 @@
+# 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 concurrent
+import sys
+import copy
+from kubernetes.client.rest import ApiException
+from sagemaker.hyperpod.cli.clients.kubernetes_client import KubernetesClient
+from sagemaker.hyperpod.cli.service.get_namespaces import GetNamespaces
+from sagemaker.hyperpod.cli.service.self_subject_access_review import SelfSubjectAccessReview
+from sagemaker.hyperpod.cli.utils import setup_logger
+
+
+logger = setup_logger(__name__)
+
+class DiscoverNamespaces:
+    def __init__(self):
+        return
+
+    def discover_accessible_namespace(self, resource_attributes_template, only_sm_managed=True):
+        """
+        Discover the accessible namespaces
+        """
+        k8s_client = KubernetesClient()
+        context_namespace = k8s_client.get_current_context_namespace()
+
+        # If the namespace is explicitly set by user, take the namespace from the config instead
+        # of discovering automatically 
+        if context_namespace is not None:
+            return context_namespace
+        
+        try:
+            if only_sm_managed:
+                namespaces = GetNamespaces().get_sagemaker_managed_namespaces()
+            else:
+                namespaces = GetNamespaces().get_namespaces()
+
+            discovered_namespaces = self.get_namespaces_by_checking_access_permission(
+                namespaces,
+                resource_attributes_template,
+            )
+
+            if len(discovered_namespaces) == 0:
+                logger.error("Found no accessible namespaces. Please ask for cluster admin for assistance or specify value of namespace explicitly in the command.")
+                sys.exit(1)
+            if len(discovered_namespaces) > 1:
+                logger.error(f"Found more than 1 accessible namespaces {discovered_namespaces}. Please specify value of namespace explicitly in the command.")
+                sys.exit(1)
+            
+            if len(discovered_namespaces) == 1:
+                logger.info(f"Found accessible namespace: {discovered_namespaces}")
+            return discovered_namespaces[0]
+        except ApiException as e:
+            if e.status == 403:
+                logger.error("Got access denied error when discovering accessible namespaces, please specify the '--namespace' parameter in the command and try again.")
+                sys.exit(1)
+
+
+    def get_namespaces_by_checking_access_permission(
+            self,
+            namespaces,
+            resource_attributes_template, 
+            max_workers=10
+        ):
+        """
+        Get the accessible namespaces by performing the SelfSubjectAccessReview. For each of the namespace,
+        check if user has specified permission to it, if the answer is NO, the namespace will be skipped.
+
+        Performing access check can take quite long if the number of namespaces is large. Thus the implementation
+        leverages the multi-threading to ensure that multiple access check can be performed in parallel.
+        """
+        subject_access_review = SelfSubjectAccessReview()
+        accessible_namespaces = list()
+        resource_attributes = list()
+
+        for namespace in namespaces:
+            resource_attribute = copy.deepcopy(resource_attributes_template)
+            resource_attribute.namespace = namespace
+            resource_attributes.append(resource_attribute)
+
+        # Multi-thread the self subject access review to improve the performance
+        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+            futures = {
+                executor.submit(
+                    subject_access_review.self_subject_access_review, resource_attribute
+                ) for resource_attribute in resource_attributes
+            }
+            
+            for future in concurrent.futures.as_completed(futures):
+                try:
+                    response = future.result()
+                    if response.status.allowed:
+                        accessible_namespaces.append(
+                            response.spec.resource_attributes.namespace
+                        )
+                except Exception as e:
+                    raise(e)
+
+        return accessible_namespaces

sagemaker/hyperpod/cli/service/exec_command.py

@@ -0,0 +1,75 @@
+# 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
+
+from sagemaker.hyperpod.cli.clients.kubernetes_client import (
+    KubernetesClient,
+)
+from sagemaker.hyperpod.cli.service.list_pods import (
+    ListPods,
+)
+
+from kubernetes.client.rest import ApiException
+
+
+class ExecCommand:
+    def __init__(self):
+        return
+
+    def exec_command(
+        self,
+        job_name: str,
+        pod_name: Optional[str],
+        namespace: Optional[str],
+        all_pods: Optional[bool],
+        bash_command: tuple,
+    ):
+        bash_command_str: str = " ".join(bash_command)
+
+        k8s_client = KubernetesClient()
+        list_pods_service = ListPods()
+
+        if not namespace:
+            namespace = k8s_client.get_current_context_namespace()
+
+        pods_for_training_job = list_pods_service.list_pods_for_training_job(
+            job_name, namespace, False
+        )
+
+        try:
+            if all_pods:
+                output = ""
+                for pod in pods_for_training_job:
+                    output += pod + "\n"
+                    output += (
+                        k8s_client.exec_command_on_pod(
+                            pod,
+                            namespace,
+                            bash_command_str,
+                        )
+                        + "\n"
+                    )
+                    output += "\n"
+                return output
+            else:
+                if pod_name not in pods_for_training_job:
+                    raise RuntimeError(
+                        f"Given pod name {pod_name} is not associated with training job {job_name} in namespace {namespace}"
+                    )
+                return k8s_client.exec_command_on_pod(
+                    pod_name,
+                    namespace,
+                    bash_command_str,
+                )
+        except ApiException as e:
+            raise RuntimeError(f"Unexpected API error: {e.reason} ({e.status})")

sagemaker/hyperpod/cli/service/get_logs.py

@@ -0,0 +1,136 @@
+# 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 sagemaker.hyperpod.cli.clients.kubernetes_client import (
+    KubernetesClient,
+)
+from sagemaker.hyperpod.cli.service.discover_namespaces import DiscoverNamespaces
+from sagemaker.hyperpod.cli.service.list_pods import (
+    ListPods,
+)
+from sagemaker.hyperpod.cli.utils import (
+    get_eks_cluster_name,
+    get_hyperpod_cluster_region,
+    validate_region_and_cluster_name,
+)
+from kubernetes.client.rest import ApiException
+from kubernetes.client import V1ResourceAttributes
+import re
+
+AMAZON_ClOUDWATCH_OBSERVABILITY = "amazon-cloudwatch-observability"
+CONTAINER_INSIGHTS_LOG_REGEX_PATTERN = "https:\/\/([a-z0-9-]+).console.aws.amazon.com\/cloudwatch\/home\?region=([a-z0-9-]+)#logsV2:log-groups\/log-group\/\$252Faws\$252Fcontainerinsights\$252F([a-zA-Z0-9-]+)\$252Fapplication\/log-events\/([a-z0-9-]+)-application.var.log.containers.([a-z0-9-]+)_([a-z0-9-]+)_([a-z0-9-]+)-([a-z0-9-]+).log"
+
+class GetLogs:
+    def __init__(self):
+        return
+
+    def get_training_job_logs(
+        self,
+        job_name: str,
+        pod_name: str,
+        namespace: Optional[str],
+    ):
+        """
+        Get logs for pod asscoited with the training job
+        """
+        k8s_client = KubernetesClient()
+        list_pods_service = ListPods()
+
+        if not namespace:
+            resource_attributes_template = V1ResourceAttributes(
+                verb="get",
+                group="",
+                resource="pods",
+                subresource="log",
+            )
+            namespace = DiscoverNamespaces().discover_accessible_namespace(
+                resource_attributes_template
+            )
+
+        try:
+            pods_for_training_job = list_pods_service.list_pods_for_training_job(
+                job_name, namespace, False
+            )
+            if pod_name not in pods_for_training_job:
+                raise RuntimeError(
+                    f"Given pod name {pod_name} is not associated with training job {job_name} in namespace {namespace}"
+                )
+            return k8s_client.get_logs_for_pod(pod_name, namespace)
+        except ApiException as e:
+            raise RuntimeError(f"Unexpected API error: {e.reason} ({e.status})")
+    
+    def generate_cloudwatch_link(
+        self,
+        pod_name: str,
+        namespace: Optional[str],
+    ):
+        eks_cluster_name = get_eks_cluster_name()
+        region = get_hyperpod_cluster_region()
+
+        if self.is_container_insights_addon_enabled(eks_cluster_name):
+            k8s_client = KubernetesClient()
+
+            # pod_details is a V1Pod object
+            pod_details = k8s_client.get_pod_details(pod_name, namespace)
+            
+            # get node name
+            if pod_details.spec and pod_details.spec.node_name:
+                node_name = pod_details.spec.node_name
+            else:
+                node_name = None
+
+            # get container name
+            if pod_details.spec and pod_details.spec.containers and pod_details.spec.containers[0].name:
+                container_name = pod_details.spec.containers[0].name
+            else:
+                container_name = None
+
+            # get container_id
+            if pod_details.status and pod_details.status.container_statuses and pod_details.status.container_statuses[0].container_id:
+                full_container_id = pod_details.status.container_statuses[0].container_id
+
+                # full_container_id has format "containerd://xxxxxxxxxx" 
+                container_id = full_container_id[13:] if full_container_id.startswith('containerd://') else None
+            else:
+                container_id = None         
+
+            # Cloudwatch container insight log groups should have the same pod log as API response
+            cloudwatch_url = self.get_log_url(eks_cluster_name, region, node_name, pod_name, namespace, container_name, container_id)
+
+            if not validate_region_and_cluster_name(region, eks_cluster_name):
+                raise ValueError('Eks cluster name or cluster region is invalid.')
+
+            if not re.match(CONTAINER_INSIGHTS_LOG_REGEX_PATTERN, cloudwatch_url):
+                raise ValueError("Failed to validate cloudwatch log url. Please make sure pod's node name, container name and container id are valid")
+            
+            cloudwatch_link = f'The pod cloudwatch log stream link is {cloudwatch_url}'
+        else:
+            cloudwatch_link = None
+
+        return cloudwatch_link
+
+    def get_log_url(self, eks_cluster_name, region, node_name, pod_name, namespace, container_name, container_id):
+        console_prefix = f'https://{region}.console.aws.amazon.com/cloudwatch/home?region={region}#'
+        log_group_prefix = f'logsV2:log-groups/log-group/$252Faws$252Fcontainerinsights$252F{eks_cluster_name}$252Fapplication/log-events/'
+        log_stream = f'{node_name}-application.var.log.containers.{pod_name}_{namespace}_{container_name}-{container_id}.log'
+        
+        return console_prefix + log_group_prefix + log_stream
+
+    def is_container_insights_addon_enabled(self, eks_cluster_name):
+        response = boto3.client("eks").list_addons(clusterName=eks_cluster_name, maxResults=50)
+        if AMAZON_ClOUDWATCH_OBSERVABILITY in response.get('addons', []):
+            return True
+        else:
+            return False