runpod-python/runpod/api/mutations/endpoints.py at 96da185b437cfcdb74f2d64cb9f75e7103cac08c · ReneNyffenegger/runpod-python · 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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
f"""Runpod | API Wrapper | Mutations | Endpoints"""

# pylint: disable=too-many-arguments


def generate_endpoint_mutation(
    name: str,
    template_id: str,
    gpu_ids: str = "AMPERE_16",
    network_volume_id: str = None,
    locations: str = None,
    idle_timeout: int = 5,
    scaler_type: str = "QUEUE_DELAY",
    scaler_value: int = 4,
    workers_min: int = 0,
    workers_max: int = 3,
    flashboot=False,
    allowed_cuda_versions: str = None,
    gpu_count: int = None,
):
    """Generate a string for a GraphQL mutation to create a new endpoint."""
    input_fields = []

    # ------------------------------ Required Fields ----------------------------- #
    if flashboot:
        input_fields.append('flashBootType: FLASHBOOT')

    input_fields.append(f'name: "{name}"')
    input_fields.append(f'templateId: "{template_id}"')
    input_fields.append(f'gpuIds: "{gpu_ids}"')

    # ------------------------------ Optional Fields ----------------------------- #
    if network_volume_id is not None:
        input_fields.append(f'networkVolumeId: "{network_volume_id}"')
    else:
        input_fields.append('networkVolumeId: ""')

    if locations is not None:
        input_fields.append(f'locations: "{locations}"')
    else:
        input_fields.append('locations: ""')

    input_fields.append(f"idleTimeout: {idle_timeout}")
    input_fields.append(f'scalerType: "{scaler_type}"')
    input_fields.append(f"scalerValue: {scaler_value}")
    input_fields.append(f"workersMin: {workers_min}")
    input_fields.append(f"workersMax: {workers_max}")

    if allowed_cuda_versions:
        input_fields.append(f'allowedCudaVersions: "{allowed_cuda_versions}"')

    if gpu_count is not None:
        input_fields.append(f"gpuCount: {gpu_count}")

    # Format the input fields into a string
    input_fields_string = ", ".join(input_fields)

    return f"""
    mutation {{
        saveEndpoint(
            input: {{
                {input_fields_string}
            }}
        ) {{
            id
            name
            templateId
            gpuIds
            networkVolumeId
            locations
            idleTimeout
            scalerType
            scalerValue
            workersMin
            workersMax
            allowedCudaVersions
            gpuCount
            flashBootType
        }}
    }}
    """


def update_endpoint_template_mutation(endpoint_id: str, template_id: str):
    """Generate a string for a GraphQL mutation to update an existing endpoint's template."""
    input_fields = []

    # ------------------------------ Required Fields ----------------------------- #
    input_fields.append(f'templateId: "{template_id}"')
    input_fields.append(f'endpointId: "{endpoint_id}"')

    # Format the input fields into a string
    input_fields_string = ", ".join(input_fields)
    result = f"""
    mutation {{
        updateEndpointTemplate(input: {{{input_fields_string}}}) {{
            id
            templateId
        }}
    }}
    """
    return result