sagemaker-hyperpod-cli/src/sagemaker/hyperpod/cli/commands/space_template.py at 11da28b3dec97257296adc279db57f2854700857 · aws/sagemaker-hyperpod-cli · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import click
import json
import yaml
from tabulate import tabulate
from sagemaker.hyperpod.space.hyperpod_space_template import HPSpaceTemplate
from sagemaker.hyperpod.common.cli_decorators import handle_cli_exceptions


@click.command("hyp-space-template")
@click.option("--file", "-f", required=True, help="YAML file containing the configuration")
@handle_cli_exceptions()
def space_template_create(file):
    """Create a space-template resource."""
    template = HPSpaceTemplate(file_path=file)
    template.create()
    click.echo(f"Space template '{template.name}' in namespace '{template.namespace}' created successfully")


@click.command("hyp-space-template")
@click.option("--namespace", "-n", required=False, default=None, help="Kubernetes namespace")
@click.option("--output", "-o", type=click.Choice(["table", "json"]), default="table")
@handle_cli_exceptions()
def space_template_list(namespace, output):
    """List space-template resources."""
    templates = HPSpaceTemplate.list(namespace)

    if output == "json":
        templates_data = [template.to_dict() for template in templates]
        click.echo(json.dumps(templates_data, indent=2))
    else:
        if templates:
            table_data = []
            for template in templates:
                table_data.append([
                    template.namespace,
                    template.name,
                    template.config_data.get("spec", {}).get("displayName", ""),
                    template.config_data.get("spec", {}).get("defaultImage", ""),
                ])
            click.echo(tabulate(table_data, headers=["NAMESPACE", "NAME", "DISPLAY_NAME", "DEFAULT_IMAGE"]))
        else:
            click.echo("No space templates found")


@click.command("hyp-space-template")
@click.option("--name", required=True, help="Name of the space template")
@click.option("--namespace", "-n", required=False, default=None, help="Kubernetes namespace")
@click.option("--output", "-o", type=click.Choice(["yaml", "json"]), default="yaml")
@handle_cli_exceptions()
def space_template_describe(name, namespace, output):
    """Describe a space-template resource."""
    template = HPSpaceTemplate.get(name, namespace)

    if output == "json":
        click.echo(json.dumps(template.to_dict(), indent=2))
    else:
        click.echo(template.to_yaml())


@click.command("hyp-space-template")
@click.option("--name", required=True, help="Name of the space template")
@click.option("--namespace", "-n", required=False, default=None, help="Kubernetes namespace")
@handle_cli_exceptions()
def space_template_delete(name, namespace):
    """Delete a space-template resource."""
    template = HPSpaceTemplate.get(name, namespace)
    template.delete()
    click.echo(f"Requested deletion for Space template '{name}' in namespace '{namespace}'")


@click.command("hyp-space-template")
@click.option("--name", required=True, help="Name of the space template")
@click.option("--namespace", "-n", required=False, default=None, help="Kubernetes namespace")
@click.option("--file", "-f", required=True, help="YAML file containing the updated template")
@handle_cli_exceptions()
def space_template_update(name, namespace, file):
    """Update a space-template resource."""
    template = HPSpaceTemplate.get(name, namespace)
    template.update(file)
    click.echo(f"Space template '{name}' in namespace '{namespace}' updated successfully")