sagemaker-hyperpod 3.0.0__py3-none-any.whl

sagemaker/hyperpod/cli/service/get_namespaces.py

@@ -0,0 +1,59 @@
+# 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 sagemaker.hyperpod.cli.clients.kubernetes_client import (
+    KubernetesClient,
+)
+
+from sagemaker.hyperpod.cli.constants.command_constants import SAGEMAKER_MANAGED_QUEUE_LABEL
+
+
+LIMIT_PER_REQUEST = 200
+
+class GetNamespaces:
+
+    def __init__(self):
+        return
+
+    def get_namespaces(self, label_selector=None):
+        """
+        Get namespaces in the cluster
+        """
+        core_v1_api = KubernetesClient().get_core_v1_api()
+        all_namespaces = list()
+        continue_token = None
+
+        while True:
+            response = None
+            if continue_token:
+                response = core_v1_api.list_namespace(
+                    limit=LIMIT_PER_REQUEST, label_selector=label_selector, _continue=continue_token
+                )
+            else:
+                response = core_v1_api.list_namespace(
+                    limit=LIMIT_PER_REQUEST, label_selector=label_selector
+                )
+
+            all_namespaces.extend([ns.metadata.name for ns in response.items])
+            continue_token = response.metadata._continue
+
+            if not continue_token:
+                break
+        
+        return all_namespaces
+
+    def get_sagemaker_managed_namespaces(self):
+        """
+        Get sagemaker managed namespaces in the cluster
+        """
+
+        return self.get_namespaces(SAGEMAKER_MANAGED_QUEUE_LABEL + "=true")

sagemaker/hyperpod/cli/service/get_training_job.py

@@ -0,0 +1,110 @@
+# 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 json
+
+from sagemaker.hyperpod.cli.clients.kubernetes_client import (
+    KubernetesClient,
+)
+
+from sagemaker.hyperpod.cli import utils
+from kubernetes.client.rest import ApiException
+from kubernetes.client import (
+    V1ResourceAttributes
+)
+
+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
+
+class GetTrainingJob:
+    def __init__(self):
+        return
+
+    def get_training_job(
+        self,
+        job_name: str,
+        namespace: Optional[str],
+        verbose: Optional[bool],
+    ):
+        """
+        Describe training job provided by the user in the specified namespace.
+        If namespace is not provided job is described from the default namespace in user context
+        """
+
+        k8s_client = KubernetesClient()
+        if not namespace:
+            resource_attributes_template = V1ResourceAttributes(
+                verb="get",
+                group=PYTORCH_CUSTOM_OBJECT_GROUP,
+                resource=PYTORCH_CUSTOM_OBJECT_PLURAL,
+            )
+            namespace = DiscoverNamespaces().discover_accessible_namespace(
+                resource_attributes_template
+            )
+
+        try:
+            result = k8s_client.get_job(
+                job_name=job_name,
+                namespace=namespace,
+            )
+        except ApiException as e:
+            raise RuntimeError(f"Unexpected API error: {e.reason} ({e.status})")
+
+        if not verbose:
+            return self._format_output_to_keep_needed_fields(result)
+        else:
+            return self._format_verbose_output(result)
+
+    def _format_output_to_keep_needed_fields(self, output):
+        result = {}
+        if output:
+            if output.get("metadata"):
+                result = {
+                    "Name": output.get("metadata").get("name"),
+                    "Namespace": output.get("metadata").get("namespace"),
+                    "Label": output.get("metadata").get("labels"),
+                    "CreationTimestamp": output.get("metadata").get(
+                        "creationTimestamp"
+                    ),
+                }
+            result.update({"Status": output.get("status")})
+            result.update({"ConsoleURL": utils.get_cluster_console_url()})
+        return json.dumps(result, indent=1, sort_keys=False)
+
+    def _format_verbose_output(self, output):
+        result = {}
+        if output:
+            if output.get("metadata"):
+                result = {
+                    "Name": output.get("metadata").get("name"),
+                    "Namespace": output.get("metadata").get("namespace"),
+                    "Label": output.get("metadata").get("labels"),
+                    "Annotations": output.get("metadata").get("annotations"),
+                    "Metadata": {
+                        "CreationTimestamp": output.get("metadata").get(
+                            "creationTimestamp"
+                        ),
+                        "Generation": output.get("metadata").get("generation"),
+                        "ResourceVersion": output.get("metadata").get(
+                            "resourceVersion"
+                        ),
+                        "UID": output.get("metadata").get("uid"),
+                    },
+                }
+            result.update({"Kind": output.get("kind")})
+            result.update({"ApiVersion": output.get("apiVersion")})
+            result.update({"Spec": output.get("spec")})
+            result.update({"Status": output.get("status")})
+            result.update({"ConsoleURL": utils.get_cluster_console_url()})
+        return json.dumps(result, indent=1, sort_keys=False)

sagemaker/hyperpod/cli/service/list_pods.py

@@ -0,0 +1,138 @@
+# 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 collections import defaultdict
+from typing import List, Optional
+import json
+
+from kubernetes.client import V1Pod, V1PodList
+
+from sagemaker.hyperpod.cli.clients.kubernetes_client import (
+    KubernetesClient,
+)
+from kubernetes.client.rest import ApiException
+from kubernetes.client import (
+    V1ResourceAttributes
+)
+
+from sagemaker.hyperpod.cli.constants.command_constants import NVIDIA_GPU_RESOURCE_LIMIT_KEY
+from sagemaker.hyperpod.cli.service.discover_namespaces import DiscoverNamespaces
+
+class ListPods:
+    def __init__(self):
+        return
+
+    def list_pods_for_training_job(
+        self,
+        job_name: str,
+        namespace: Optional[str],
+        pretty: Optional[bool],
+    ):
+        """
+        List pods associated with a training job
+        """
+        k8s_client = KubernetesClient()
+
+        if not namespace:
+            resource_attributes_template = V1ResourceAttributes(
+                verb="list",
+                group="",
+                resource="pods",
+            )
+            namespace = DiscoverNamespaces().discover_accessible_namespace(
+                resource_attributes_template
+            )
+
+        label_filter = f"training.kubeflow.org/job-name={job_name}"
+
+        try:
+            _pods: V1PodList = k8s_client.list_pods_with_labels(namespace, label_filter)
+        except ApiException as e:
+            raise RuntimeError(f"Unexpected API error: {e.reason} ({e.status})")
+
+        if pretty:
+            return self._generate_list_pods_output(_pods)
+        else:
+            return self._generate_pods_list(_pods)
+
+    def list_pods_and_get_requested_resources_group_by_node_name(
+        self,
+    ):
+        """
+        List pods for all namespaces and initialized by kubeflow.
+        Group by the node_name of the pod and the value is the
+        accelerator devices resources requested
+        """
+        k8s_client = KubernetesClient()
+
+        label_filter = f"training.kubeflow.org/job-name"
+        pods = k8s_client.list_pods_in_all_namespaces_with_labels(label_filter)
+
+        # Dictionary to hold total GPU/Neuron requests per node
+        accelerator_devices_requests_by_node = defaultdict(int)
+
+        # Loop through each pod and aggregate the GPU resource requests per node
+        for pod in pods:
+            node_name = pod.spec.node_name
+            # Check if the pod has a node_name assigned, if pod is not scheduled, it may not have node_name
+            if node_name:
+                for container in pod.spec.containers:
+                    if container.resources and container.resources.requests:
+                        gpu_request = container.resources.requests.get(NVIDIA_GPU_RESOURCE_LIMIT_KEY)
+                        neuron_request = container.resources.requests.get(
+                            "aws.amazon.com/neurondevice"
+                        )
+                        if gpu_request:
+                            accelerator_devices_requests_by_node[node_name] += int(
+                                gpu_request
+                            )
+                        if neuron_request:
+                            accelerator_devices_requests_by_node[node_name] += int(
+                                neuron_request
+                            )
+
+        return accelerator_devices_requests_by_node
+
+    def _generate_list_pods_output(self, pods: V1PodList) -> Optional[str]:
+        output_pods = {"pods": []}
+        if pods.items and len(pods.items) > 0:
+            _pod: V1Pod
+            for _pod in pods.items:
+                if _pod.metadata and _pod.metadata.name and _pod.metadata.namespace:
+                    name = _pod.metadata.name
+                    namespace = _pod.metadata.namespace
+                    status = None
+                    creation_timestamp = None
+                    if _pod.status and _pod.status.phase:
+                        status = _pod.status.phase
+                    if _pod.metadata.creation_timestamp:
+                        creation_timestamp = str(_pod.metadata.creation_timestamp)
+                    output_pods["pods"].append(
+                        {
+                            "PodName": name,
+                            "Namespace": namespace,
+                            "Status": status,
+                            "CreationTime": creation_timestamp,
+                        }
+                    )
+
+        return json.dumps(output_pods, indent=1, sort_keys=False)
+
+    def _generate_pods_list(self, pods: V1PodList) -> List:
+        output_pods = []
+        if pods.items and len(pods.items) > 0:
+            _pod: V1Pod
+            for _pod in pods.items:
+                if _pod.metadata and _pod.metadata.name:
+                    name = _pod.metadata.name
+                    output_pods.append(name)
+        return output_pods

sagemaker/hyperpod/cli/service/list_training_jobs.py

@@ -0,0 +1,178 @@
+# 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 List, Optional
+
+import json
+from datetime import datetime
+
+from sagemaker.hyperpod.cli.clients.kubernetes_client import (
+    KubernetesClient,
+)
+from sagemaker.hyperpod.cli.utils import setup_logger
+from kubernetes.client.rest import ApiException
+from kubernetes.client import (
+    V1ResourceAttributes
+)
+
+from sagemaker.hyperpod.cli.constants.command_constants import KUEUE_WORKLOAD_PRIORITY_CLASS_LABEL_KEY, OutputFormat
+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 tabulate import tabulate
+
+
+class ListTrainingJobs:
+    def __init__(self):
+        return
+
+    def list_training_jobs(
+        self,
+        namespace: Optional[str],
+        all_namespaces: Optional[bool],
+        selector: Optional[str],
+        output: Optional[str],
+    ) -> str:
+        """
+        List training job provided by the user in the specified namespace.
+        If namespace is not provided job are listed the default namespace in user context
+        If all_namespace is true we will list training job from all namespaces that user has access
+        Selector when specified will filter list of job addisional based on labels filter provided
+        """
+        k8s_client = KubernetesClient()
+
+        jobs: List = []
+        logger = setup_logger(__name__)
+        logger.debug(namespace)
+        try:
+            if all_namespaces:
+                namespaces: List[str] = k8s_client.list_namespaces()
+                for _namespace in namespaces:
+                    _jobs = k8s_client.list_training_jobs(
+                        namespace=_namespace,
+                        label_selector=selector,
+                    )
+                    if _jobs.get("items") and len(_jobs.get("items")) > 0:
+                        for _job in _jobs.get("items"):
+                            jobs.append(_job)
+            else:
+                if not namespace:
+                    resource_attributes_template = V1ResourceAttributes(
+                        verb="list",
+                        group=PYTORCH_CUSTOM_OBJECT_GROUP,
+                        resource=PYTORCH_CUSTOM_OBJECT_PLURAL,
+                    )
+                    namespace = DiscoverNamespaces().discover_accessible_namespace(
+                        resource_attributes_template
+                    )
+                else:
+                    if not k8s_client.check_if_namespace_exists(namespace):
+                        raise ValueError(f"Namespace {namespace} does not exist!")
+
+                namespace_jobs = k8s_client.list_training_jobs(
+                    namespace=namespace,
+                    label_selector=selector,
+                )
+                if namespace_jobs.get("items") and len(namespace_jobs.get("items")) > 0:
+                    for namespace_job in namespace_jobs.get("items"):
+                        jobs.append(namespace_job)
+        except ApiException as e:
+            raise RuntimeError(f"Unexpected API error: {e.reason} ({e.status})")
+
+        return self._generate_list_training_job_output(jobs, output)
+
+    def _generate_list_training_job_output(self, jobs: List, output: Optional[str]):
+        output_jobs = {"jobs": []}
+        priority_header_required = False
+
+        for job in jobs:
+            if job.get("metadata"):
+                name = job.get("metadata").get("name")
+                namespace = job.get("metadata").get("namespace")
+                creation_time = None
+                state = None
+                priority = self._get_job_priority(job)
+                if job.get("status"):
+                    creation_time = job.get("status").get("startTime")
+                    state = self._get_job_status(job.get("status").get("conditions"))
+
+                job_summary = {
+                    "Name": name,
+                    "Namespace": namespace,
+                    "CreationTime": creation_time,
+                    "State": state,
+                }
+
+                if priority is not None:
+                    job_summary["priority"] = priority
+                    priority_header_required = True
+
+                output_jobs["jobs"].append(job_summary)
+
+        if output == OutputFormat.TABLE.value:
+            return self._generate_table(output_jobs, priority_header_required)
+        return json.dumps(output_jobs, indent=4, sort_keys=False)
+
+    def _generate_table(self, output_jobs, priority_header_required):
+        headers = [
+                "Name",
+                "Namespace",
+                "CreationTime",
+                "State"
+            ]
+
+        if priority_header_required:
+            headers.append("Priority")
+
+        jobs = []
+        if "jobs" in output_jobs and isinstance(output_jobs["jobs"], list):
+            for job in output_jobs["jobs"]:
+                job_values = list(job.values())
+                if priority_header_required and len(job_values) == 4:
+                    job_values.append("NA")
+                jobs.append(job_values)
+
+        return tabulate(jobs, headers=headers, tablefmt="presto")
+
+    def _get_job_status(self, status: List) -> Optional[str]:
+        current_status = None
+        last_date_time = datetime.strptime(
+            "2001-08-27T22:47:57Z",
+            "%Y-%m-%dT%H:%M:%SZ",
+        )
+        for state in status:
+            state_date_time = datetime.strptime(
+                state.get("lastTransitionTime"),
+                "%Y-%m-%dT%H:%M:%SZ",
+            )
+            if state_date_time >= last_date_time:
+                last_date_time = state_date_time
+                current_status = state.get("type")
+        return current_status
+
+    def _get_job_priority(self, job) -> Optional[str]:
+        worker_template = job.get("spec", {}).get("pytorchReplicaSpecs", {}).get("Worker", {}).get("template", {})
+        wl_priority_cls = worker_template.get("metadata", {}).get("labels", {}).get(KUEUE_WORKLOAD_PRIORITY_CLASS_LABEL_KEY, None)
+        spec_priority_cls = worker_template.get("spec", {}).get("priorityClassName", None)
+
+        # For reference: https://kueue.sigs.k8s.io/docs/concepts/workload_priority_class/
+        # The workload priority class takes precedence over the k8s priority class.
+        # Because the cli focuses on the job level which means workload priority is more essential.
+        # There is possibility that these two priorities coexist at the same time. In this case,
+        # the k8s priority class will be used as pod priority. cli should still take workload
+        # priority in this scenario.
+
+        if wl_priority_cls is not None:
+            return wl_priority_cls
+        elif spec_priority_cls is not None:
+            return spec_priority_cls
+        
+        return None

sagemaker/hyperpod/cli/service/self_subject_access_review.py

@@ -0,0 +1,47 @@
+# 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 sagemaker.hyperpod.cli.clients.kubernetes_client import (
+    KubernetesClient,
+)
+from kubernetes.client import (
+    V1SelfSubjectAccessReview, 
+    V1SelfSubjectAccessReviewSpec, 
+)
+
+
+class SelfSubjectAccessReview:
+    def __init__(self):
+        return
+
+    def self_subject_access_review(
+            self, 
+            resource_attributes=None, 
+            non_resource_attributes=None, 
+            local_vars_configuration=None,
+        ):
+        """
+        Submit self subject access review
+        """
+        auth_v1_api = KubernetesClient().get_auth_v1_api()
+
+        access_review = V1SelfSubjectAccessReview(
+            spec=V1SelfSubjectAccessReviewSpec(
+                resource_attributes=resource_attributes,
+                non_resource_attributes=non_resource_attributes,
+                local_vars_configuration=local_vars_configuration,
+            )
+        )
+
+        response = auth_v1_api.create_self_subject_access_review(body=access_review)
+        
+        return response

sagemaker/hyperpod/cli/telemetry/__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/telemetry/user_agent.py

@@ -0,0 +1,36 @@
+# 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.
+"""Placeholder docstring"""
+
+from __future__ import absolute_import
+
+import importlib.metadata
+
+
+CLI_PREFIX = "AWS-SageMaker-Hyperpod-CLI"
+
+
+def get_user_agent_extra_suffix():
+    """Get the user agent extra suffix string specific to SageMaker Hyperpod CLI
+
+    Adhers to new boto recommended User-Agent 2.0 header format
+
+    Returns:
+        str: The user agent extra suffix string to be appended
+    """
+    suffix = "cli/{}#{}".format(
+        CLI_PREFIX,
+        importlib.metadata.version("sagemaker-hyperpod"),
+    )
+
+    return suffix

sagemaker/hyperpod/cli/templates/__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/templates/k8s_pytorch_job_template.py

@@ -0,0 +1,74 @@
+# 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.
+KUBERNETES_PYTORCH_JOB_TEMPLATE = """### Please keep template file unchanged ###
+defaults:
+    - override hydra/job_logging: stdout
+
+hydra:
+    run:
+        dir: .
+    output_subdir: null
+
+training_cfg:
+    entry_script: ??? # Path to the entry script of training/fine-tuning. This path should be inside container or relative path in git repo
+    script_args: ??? # Entry script arguments
+    run:
+        nodes: ??? # Number of nodes to use for current training
+        ntasks_per_node: ??? # Number of tasks per node
+cluster:
+    cluster_type: k8s  # currently k8s only
+    instance_type: ???
+    cluster_config:
+        namespace: ??? # the namespace to submit job
+        custom_labels: ???
+        service_account_name: null
+        annotations: ???
+        priority_class_name: ???
+        # Create k8s NodeAffinity to select nodes to deploy jobs which matches required and preferred labels
+        # Structure:
+        #   label_selector:
+        #     required: <required label key-values pair>
+        #     preferred: <preferred label key-values pair>
+        #     weights: <weights list used by preferred labels to get nodes priority>
+        # Example:
+        #   label_selector:
+        #     required:
+        #       example-label-key:
+        #         - expected-label-value-1
+        #         - expected-label-value-2
+        #     preferred:
+        #       preferred-label-key:
+        #         - preferred-label-value-1
+        #         - preferred-label-value-2
+        #     weights:
+        #       - 100
+        label_selector: ???
+        # persistent volume, usually used to mount FSx
+        persistent_volume_claims: null
+        pullPolicy: ??? # policy to pull container, can be Always, IfNotPresent and Never
+        restartPolicy: ??? # PyTorchJob restart policy
+        # temp volume, usually used to mount temp directory
+        # volumes, used to mount temp path to container
+        # example:
+        # volumes:
+        #  - volumeName: data1
+        #    hostPath: "/data"
+        #    mountPath: "/data"              
+        volumes: null
+
+base_results_dir: ???  # Location to store the results, checkpoints and logs.
+container: ??? # container to use
+
+env_vars:
+    NCCL_DEBUG: INFO # Logging level for NCCL. Set to "INFO" for debug information
+"""