Skip to content

Commit

Permalink
Add ComputeClientV2 methods to manage endpoints
Browse files Browse the repository at this point in the history
Added `.register_endpoint()`, `.get_endpoint()`, `.get_endpoint_status()`
`.get_endpoints()`, `.delete_endpoint()`, and `.lock_endpoint()` methods
to the `ComputeClientV2` class.
  • Loading branch information
rjmello committed Nov 19, 2024
1 parent 81446cd commit 874682f
Show file tree
Hide file tree
Showing 14 changed files with 331 additions and 0 deletions.
5 changes: 5 additions & 0 deletions src/globus_sdk/_testing/data/compute/_common.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,11 @@
import uuid

USER_ID = str(uuid.uuid1())

SUBSCRIPTION_ID = str(uuid.uuid1())

ENDPOINT_ID = str(uuid.uuid1())
ENDPOINT_ID_2 = str(uuid.uuid1())

FUNCTION_ID = str(uuid.uuid1())
FUNCTION_NAME = "howdy_world"
Expand Down
13 changes: 13 additions & 0 deletions src/globus_sdk/_testing/data/compute/v2/delete_endpoint.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
from globus_sdk._testing.models import RegisteredResponse, ResponseSet

from .._common import ENDPOINT_ID

RESPONSES = ResponseSet(
metadata={"endpoint_id": ENDPOINT_ID},
default=RegisteredResponse(
service="compute",
path=f"/v2/endpoints/{ENDPOINT_ID}",
method="DELETE",
json={"result": 302},
),
)
37 changes: 37 additions & 0 deletions src/globus_sdk/_testing/data/compute/v2/get_endpoint.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
from globus_sdk._testing.models import RegisteredResponse, ResponseSet

from .._common import ENDPOINT_ID, SUBSCRIPTION_ID

ENDPOINT_CONFIG = """
display_name: My Endpoint
engine:
type: GlobusComputeEngine
"""

DEFAULT_RESPONSE_DOC = {
"uuid": ENDPOINT_ID,
"name": "my-endpoint",
"display_name": "My Endpoint",
"multi_user": False,
"public": False,
"endpoint_config": ENDPOINT_CONFIG.strip(),
"user_config_template": "",
"user_config_schema": {},
"description": "My endpoint description",
"hostname": "my-hostname",
"local_user": "user1",
"ip_address": "140.221.112.13",
"endpoint_version": "2.31.0",
"sdk_version": "2.31.0",
"subscription_uuid": SUBSCRIPTION_ID,
}

RESPONSES = ResponseSet(
metadata={"endpoint_id": ENDPOINT_ID},
default=RegisteredResponse(
service="compute",
path=f"/v2/endpoints/{ENDPOINT_ID}",
method="GET",
json=DEFAULT_RESPONSE_DOC,
),
)
24 changes: 24 additions & 0 deletions src/globus_sdk/_testing/data/compute/v2/get_endpoint_status.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
from globus_sdk._testing.models import RegisteredResponse, ResponseSet

from .._common import ENDPOINT_ID

DEFAULT_RESPONSE_DOC = {
"status": "online",
"details": {
"total_workers": 1,
"idle_workers": 0,
"pending_tasks": 0,
"outstanding_tasks": 0,
"managers": 1,
},
}

RESPONSES = ResponseSet(
metadata={"endpoint_id": ENDPOINT_ID},
default=RegisteredResponse(
service="compute",
path=f"/v2/endpoints/{ENDPOINT_ID}/status",
method="GET",
json=DEFAULT_RESPONSE_DOC,
),
)
28 changes: 28 additions & 0 deletions src/globus_sdk/_testing/data/compute/v2/get_endpoints.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
from globus_sdk._testing.models import RegisteredResponse, ResponseSet

from .._common import ENDPOINT_ID, ENDPOINT_ID_2, USER_ID

DEFAULT_RESPONSE_DOC = [
{
"uuid": ENDPOINT_ID,
"name": "my-endpoint",
"display_name": "My Endpoint",
"owner": USER_ID,
},
{
"uuid": ENDPOINT_ID_2,
"name": "my-second-endpoint",
"display_name": "My Second Endpoint",
"owner": USER_ID,
},
]

RESPONSES = ResponseSet(
metadata={"endpoint_id": ENDPOINT_ID, "endpoint_id_2": ENDPOINT_ID_2},
default=RegisteredResponse(
service="compute",
path="/v2/endpoints",
method="GET",
json=DEFAULT_RESPONSE_DOC,
),
)
18 changes: 18 additions & 0 deletions src/globus_sdk/_testing/data/compute/v2/lock_endpoint.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
from globus_sdk._testing.models import RegisteredResponse, ResponseSet

from .._common import ENDPOINT_ID

DEFAULT_RESPONSE_DOC = {
"endpoint_id": ENDPOINT_ID,
"lock_expiration_timestamp": "2021-07-01T00:00:00.000000",
}

RESPONSES = ResponseSet(
metadata={"endpoint_id": ENDPOINT_ID},
default=RegisteredResponse(
service="compute",
path=f"/v2/endpoints/{ENDPOINT_ID}/lock",
method="POST",
json=DEFAULT_RESPONSE_DOC,
),
)
29 changes: 29 additions & 0 deletions src/globus_sdk/_testing/data/compute/v2/register_endpoint.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
from globus_sdk._testing.models import RegisteredResponse, ResponseSet

from .._common import ENDPOINT_ID

DEFAULT_RESPONSE_DOC = {
"endpoint_id": ENDPOINT_ID,
"task_queue_info": {
"connection_url": "amqps://user:[email protected]",
"exchange": "some_exchange",
"queue": "some_queue",
},
"result_queue_info": {
"connection_url": "amqps://user:[email protected]",
"exchange": "some_exchange",
"queue": "some_queue",
"queue_publish_kwargs": {},
},
"warnings": [],
}

RESPONSES = ResponseSet(
metadata={"endpoint_id": ENDPOINT_ID},
default=RegisteredResponse(
service="compute",
path="/v2/endpoints",
method="POST",
json=DEFAULT_RESPONSE_DOC,
),
)
88 changes: 88 additions & 0 deletions src/globus_sdk/services/compute/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,94 @@ class ComputeClientV2(client.BaseClient):
scopes = ComputeScopes
default_scope_requirements = [Scope(ComputeScopes.all)]

def register_endpoint(self, data: dict[str, t.Any]) -> GlobusHTTPResponse:
"""Register a new endpoint.
:param data: An endpoint registration document.
.. tab-set::
.. tab-item:: API Info
.. extdoclink:: Register Endpoint
:service: compute
:ref: Endpoints/operation/register_endpoint_v2_endpoints_post
"""
return self.post("/v2/endpoints", data=data)

def get_endpoint(self, endpoint_id: UUIDLike) -> GlobusHTTPResponse:
"""Get information about a registered endpoint.
:param endpoint_id: The ID of the Globus Compute endpoint.
.. tab-set::
.. tab-item:: API Info
.. extdoclink:: Get Endpoint
:service: compute
:ref: Endpoints/operation/get_endpoint_v2_endpoints__endpoint_uuid__get
""" # noqa: E501
return self.get(f"/v2/endpoints/{endpoint_id}")

def get_endpoint_status(self, endpoint_id: UUIDLike) -> GlobusHTTPResponse:
"""Get the status of a registered endpoint.
:param endpoint_id: The ID of the Globus Compute endpoint.
.. tab-set::
.. tab-item:: API Info
.. extdoclink:: Get Endpoint Status
:service: compute
:ref: Endpoints/operation/get_endpoint_status_v2_endpoints__endpoint_uuid__status_get
""" # noqa: E501
return self.get(f"/v2/endpoints/{endpoint_id}/status")

def get_endpoints(self) -> GlobusHTTPResponse:
"""Get a list of registered endpoints associated with the authenticated user.
.. tab-set::
.. tab-item:: API Info
.. extdoclink:: Get Endpoints
:service: compute
:ref: Endpoints/operation/get_endpoints_v2_endpoints_get
""" # noqa: E501
return self.get("/v2/endpoints")

def delete_endpoint(self, endpoint_id: UUIDLike) -> GlobusHTTPResponse:
"""Delete a registered endpoint.
:param endpoint_id: The ID of the Globus Compute endpoint.
.. tab-set::
.. tab-item:: API Info
.. extdoclink:: Delete Endpoint
:service: compute
:ref: Endpoints/operation/delete_endpoint_v2_endpoints__endpoint_uuid__delete
""" # noqa: E501
return self.delete(f"/v2/endpoints/{endpoint_id}")

def lock_endpoint(self, endpoint_id: UUIDLike) -> GlobusHTTPResponse:
"""Temporarily block registration requests for the endpoint.
:param endpoint_id: The ID of the Globus Compute endpoint.
.. tab-set::
.. tab-item:: API Info
.. extdoclink:: Delete Endpoint
:service: compute
:ref: Endpoints/operation/lock_endpoint_v2_endpoints__endpoint_uuid__lock_post
""" # noqa: E501
return self.post(f"/v2/endpoints/{endpoint_id}/lock")

def register_function(
self,
function_data: dict[str, t.Any],
Expand Down
9 changes: 9 additions & 0 deletions tests/functional/services/compute/v2/test_delete_endpoint.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
import globus_sdk
from globus_sdk._testing import load_response


def test_delete_endpoint(compute_client_v2: globus_sdk.ComputeClientV2):
meta = load_response(compute_client_v2.delete_endpoint).metadata
res = compute_client_v2.delete_endpoint(endpoint_id=meta["endpoint_id"])
assert res.http_status == 200
assert res.data == {"result": 302}
9 changes: 9 additions & 0 deletions tests/functional/services/compute/v2/test_get_endpoint.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
import globus_sdk
from globus_sdk._testing import load_response


def test_get_endpoint(compute_client_v2: globus_sdk.ComputeClientV2):
meta = load_response(compute_client_v2.get_endpoint).metadata
res = compute_client_v2.get_endpoint(endpoint_id=meta["endpoint_id"])
assert res.http_status == 200
assert res.data["uuid"] == meta["endpoint_id"]
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
import globus_sdk
from globus_sdk._testing import load_response


def test_get_endpoint_status(compute_client_v2: globus_sdk.ComputeClientV2):
meta = load_response(compute_client_v2.get_endpoint_status).metadata
res = compute_client_v2.get_endpoint_status(endpoint_id=meta["endpoint_id"])
assert res.http_status == 200
assert res.data["status"] == "online"
12 changes: 12 additions & 0 deletions tests/functional/services/compute/v2/test_get_endpoints.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
import globus_sdk
from globus_sdk._testing import load_response


def test_get_endpoints(compute_client_v2: globus_sdk.ComputeClientV2):
meta = load_response(compute_client_v2.get_endpoints).metadata

res = compute_client_v2.get_endpoints()

assert res.http_status == 200
assert res.data[0]["uuid"] == meta["endpoint_id"]
assert res.data[1]["uuid"] == meta["endpoint_id_2"]
11 changes: 11 additions & 0 deletions tests/functional/services/compute/v2/test_lock_endpoint.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
import globus_sdk
from globus_sdk._testing import load_response


def test_lock_endpoint(compute_client_v2: globus_sdk.ComputeClientV2):
meta = load_response(compute_client_v2.lock_endpoint).metadata

res = compute_client_v2.lock_endpoint(endpoint_id=meta["endpoint_id"])

assert res.http_status == 200
assert res.data["endpoint_id"] == meta["endpoint_id"]
39 changes: 39 additions & 0 deletions tests/functional/services/compute/v2/test_register_endpoint.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
import uuid

import globus_sdk
from globus_sdk._testing import load_response

ENDPOINT_CONFIG = """
display_name: My Endpoint
engine:
type: GlobusComputeEngine
"""


def test_register_endpoint(compute_client_v2: globus_sdk.ComputeClientV2):
meta = load_response(compute_client_v2.register_endpoint).metadata
register_doc = {
"endpoint_uuid": meta["endpoint_id"],
"endpoint_name": "my-endpoint",
"display_name": "My Endpoint",
"version": "2.31.0",
"multi_user": False,
"allowed_functions": [str(uuid.uuid1())],
"authentication_policy": str(uuid.uuid1()),
"metadata": {
"endpoint_config": ENDPOINT_CONFIG.strip(),
"user_config_template": "",
"user_config_schema": {},
"description": "My endpoint description",
"ip_address": "140.221.112.13",
"hostname": "my-hostname",
"local_user": "user1",
"sdk_version": "2.31.0",
"endpoint_version": "2.31.0",
},
}

res = compute_client_v2.register_endpoint(data=register_doc)

assert res.http_status == 200
assert res.data["endpoint_id"] == meta["endpoint_id"]

0 comments on commit 874682f

Please sign in to comment.