Merge pull request #100 from NERSC/globus

Author: tylern4

pyproject.toml

@@ -14,7 +14,7 @@ authors = [
 ]
 description = "Python client for NERSC SF API"
 readme = "README.md"
-requires-python = ">=3.7"
+requires-python = ">=3.8"
 dependencies = [
   "authlib",
   "httpx",

scripts/run.py

@@ -34,6 +34,9 @@ def run_unasync():
         "AsyncJobSacct": "JobSacct",
         "AsyncJobSqueue": "JobSqueue",
         "AsyncOAuth2Client": "OAuth2Client",
+        "AsyncStorage": "Storage",
+        "AsyncGlobusStorage": "GlobusStorage",
+        "AsyncGlobusTransfer": "GlobusTransfer",
         "aclose": "close",
         "_ASYNC_SLEEP": "_SLEEP",
     }
@@ -100,7 +103,7 @@ def _to_str_enum(code: str) -> str:
 def _fix_date_import(code: str) -> str:
     replacements = {
         "import date": "import date as date_",
-        "date: Optional\[date\]": "date: Optional[date_]",
+        "date: Optional\\[date\\]": "date: Optional[date_]",
     }
 
     for target, replacement in replacements.items():
@@ -193,7 +196,7 @@ def resources_codegen(
         name = name if name != "int" else f"_{name}"
         name = name.replace("-", "_")
         names.append(f'    {name} = "{value}"')
-        names.append(f"    \"\"\" {s['system_type']}: {s['full_name']}\"\"\"")
+        names.append(f'    """ {s["system_type"]}: {s["full_name"]}"""')
 
     resources = "\n".join(names)
     now = datetime.datetime.now()

src/sfapi_client/_async/client.py

@@ -21,6 +21,7 @@
 from .._models.resources import Resource
 from .groups import AsyncGroup, AsyncGroupMember
 from .users import AsyncUser
+from .storage import AsyncGlobusStorage, AsyncStorage
 from .projects import AsyncProject, AsyncRole
 from .paths import AsyncRemotePath
 
@@ -252,6 +253,7 @@ def __init__(
         self.__http_client = None
         self._api = None
         self._resources = None
+        self._storage = None
         self._wait_interval = wait_interval
         self._access_token = access_token
 
@@ -504,15 +506,27 @@ def resources(self) -> AsyncResources:
 
         return self._resources
 
+    @property
+    def storage(self) -> AsyncStorage:
+        """Storage related objects and methods
+
+        :return AsyncStorage: Object with storage related objects and methods
+        """
+        if self._storage is None:
+            self._storage = AsyncStorage(self)
+        return self._storage
+
 
 AsyncCompute.model_rebuild()
+AsyncGlobusStorage.model_rebuild()
 AsyncGroup.model_rebuild()
 AsyncUser.model_rebuild()
 AsyncProject.model_rebuild()
 AsyncRemotePath.model_rebuild()
 AsyncRole.model_rebuild()
 AsyncGroupMember.model_rebuild()
 
+
 # Ensure that the job models are built, we need to import here to
 # avoid circular imports
 from .jobs import AsyncJobSacct, AsyncJobSqueue  # noqa: E402

src/sfapi_client/_async/compute.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional, Union, Callable
+from typing import Dict, List, Optional, Union
 import json
 from pydantic import PrivateAttr, ConfigDict
 from ..exceptions import SfApiError
@@ -7,7 +7,6 @@
 from .._models import (
     AppRoutersStatusModelsStatus as ComputeBase,
     Task,
-    StatusValue,
     PublicHost as Machine,
     BodyRunCommandUtilitiesCommandMachinePost as RunCommandBody,
     AppRoutersComputeModelsCommandOutput as RunCommandResponse,
@@ -16,26 +15,12 @@
 from .paths import AsyncRemotePath
 from .._monitor import AsyncJobMonitor
 from .._compute import CommandResult, SubmitJobResponse, SubmitJobResponseStatus
+from .._utils import check_auth
 
 # Patch to return str names from Enum of py3.11
 Machine.__str__ = lambda self: self.value
 
 
-def check_auth(method: Callable):
-    def wrapper(self, *args, **kwargs):
-        if self.client._client_id is None and self.client._access_token is None:
-            raise SfApiError(
-                f"Cannot call {self.__class__.__name__}.{method.__name__}() with an unauthenticated client."  # noqa: E501
-            )
-        elif self.status in [StatusValue.unavailable]:
-            raise SfApiError(
-                f"Compute resource {self.name} is {self.status.name}, {self.notes}"
-            )
-        return method(self, *args, **kwargs)
-
-    return wrapper
-
-
 class AsyncCompute(ComputeBase):
     client: Optional["AsyncClient"]  # noqa: F821
     _monitor: AsyncJobMonitor = PrivateAttr()

src/sfapi_client/_async/groups.py

@@ -1,21 +1,9 @@
-from typing import Optional, Union, List, Any, Callable
-from functools import wraps
+from typing import Optional, Union, List, Any
 from pydantic import ValidationError, Field, BaseModel, ConfigDict
 from .._models import BatchGroupAction as GroupAction, UserStats as GroupMemberBase
 from ..exceptions import SfApiError
 from .users import AsyncUser
-
-
-def check_auth(method: Callable):
-    @wraps(method)
-    def wrapper(self, *args, **kwargs):
-        if self._client_id is None:
-            raise SfApiError(
-                f"Cannot call {self.__class__.__name__}.{method.__name__}() with an unauthenticated client."  # noqa: E501
-            )
-        return method(self, *args, **kwargs)
-
-    return wrapper
+from .._utils import check_auth
 
 
 class AsyncGroupMember(GroupMemberBase):

src/sfapi_client/_async/storage.py

@@ -0,0 +1,210 @@
+from pathlib import Path
+from typing import Optional, Union, List, Any
+from pydantic import ConfigDict
+import sys
+import math
+
+from .._utils import _ASYNC_SLEEP
+from ..paths import AsyncRemotePath
+from .._utils import check_auth
+
+from .compute import Machine
+from .._models import (
+    AppRoutersStatusModelsStatus as StorageBase,
+    GlobusTransfer as GlobusTransferModel,
+    BodyStartGlobusTransferStorageGlobusTransferPost as GlobusBodyPost,
+    GlobusTransferResult,
+    GlobusStatus,
+)
+
+GLOBUS_TERMINAL_STATES = [
+    GlobusStatus.CANCELED,
+    GlobusStatus.FAILED,
+    GlobusStatus.SUCCEEDED,
+]
+
+
+class AsyncStorage:
+    def __init__(self, client: "AsyncClient"):  # noqa: F821
+        self.client = client
+
+    async def globus(
+        self,
+    ):
+        """Create a globus transfer object to start and monitor transfers
+
+        - Must select the Globus option when creating the SuperFacility key
+
+        ```python
+        >>> from sfapi_client import AsyncClient
+        >>> async with AsyncClient(client_id, client_secret) as client:
+        >>>     globus = client.storage.globus()
+        ```
+
+        :return AsyncGlobusStorage: Globus object to start and monitor transfers
+        """
+        response = await self.client.get("status/globus")
+        values = response.json()
+        values["client"] = self.client
+        _globus = AsyncGlobusStorage.model_validate(values)
+
+        return _globus
+
+
+class AsyncGlobusTransfer(GlobusTransferResult):
+    globus: Optional["AsyncGlobusStorage"]  # noqa: F821
+    transfer_id: str
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    async def update(self):
+        """Updates the status of the transfer"""
+        job_state = await self._fetch_state()
+        self._update(job_state)
+
+    def _update(self, new_job_state: Any):
+        for k in new_job_state.model_fields_set:
+            v = getattr(new_job_state, k)
+            setattr(self, k, v)
+
+        return self
+
+    async def _wait_until(self, states: List[GlobusStatus], timeout: int = sys.maxsize):
+        max_iteration = math.ceil(timeout / self.globus.client._wait_interval)
+        iteration = 0
+
+        while self.globus_status not in states:
+            await self.update()
+            await _ASYNC_SLEEP(self.globus.client._wait_interval)
+
+            if iteration == max_iteration:
+                raise TimeoutError()
+
+            iteration += 1
+
+        return self.globus_status
+
+    async def _wait_until_complete(self, timeout: int = sys.maxsize):
+        return await self._wait_until(GLOBUS_TERMINAL_STATES, timeout)
+
+    def __await__(self):
+        return self._wait_until_complete().__await__()
+
+    async def complete(self, timeout: int = sys.maxsize):
+        """Wait for the transfer to complete
+
+        >>> from sfapi_client import AsyncClient
+        >>> async with AsyncClient(client_id, client_secret) as client:
+        >>>     globus = client.storage.globus()
+        >>>     res = await globus.transfer(
+                        "globus-transfer-uuid"
+                )
+        >>>     await res.complete()
+
+        :param int timeout: time to wait for the transfer to complete, defaults to sys.maxsize
+        :return GlobusStart: Gives the file status for the transfer
+        """
+        return await self._wait_until_complete(timeout)
+
+    async def _fetch_state(self):
+        r = await self.globus.client.get(f"storage/globus/transfer/{self.transfer_id}")
+        json_response = r.json()
+        json_response["transfer_id"] = self.transfer_id
+        json_response["globus"] = self.globus
+        transfer = AsyncGlobusTransfer.model_validate(json_response)
+        return transfer
+
+
+class AsyncGlobusStorage(StorageBase):
+    client: Optional["AsyncClient"]  # noqa: F821
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    @check_auth
+    async def start_transfer(
+        self,
+        source_machine: Union[Machine, str],
+        target_machine: Union[Machine, str],
+        source_dir: Union[str, Path, AsyncRemotePath],
+        target_dir: Union[str, Path, AsyncRemotePath],
+        label: Optional[str] = None,
+    ) -> AsyncGlobusTransfer:
+        """Start a Globus transfer throught the SuperFacility API
+
+        - Must select the Globus option when creating the SuperFacility key
+
+        ```python
+        >>> from sfapi_client import AsyncClient
+        >>> async with AsyncClient(client_id, client_secret) as client:
+        >>>     globus_client = client.storage.globus()
+        >>>     res = await globus_client.start_transfer(
+                        Machine.Perlmutter,
+                        "/pscratch/sd/u/user/globus",
+                        Machine.dtns,
+                        "/global/cfs/cdirs/m0000/globus"
+                    )
+        ```
+
+        :param str source_dir: Path to file or directory on the source to transfer
+        :param str target_dir: Path to directory on the target to transfer files to
+        :param Optional[str] label: Label for the transfer,
+        defaults to None and the API will create a label for the transfer
+        :return AsyncGlobusTransfer
+        """
+
+        if None in [source_machine, source_dir, target_machine, target_dir]:
+            # Check that all parametes are not none
+            raise ValueError("sources, and targets cannot be None")
+
+        # Make machine names match those in the API endpoint
+        source_name = (
+            "dtn" if source_machine in [Machine.dtns, Machine.dtn01] else source_machine
+        )
+        target_name = (
+            "dtn" if target_machine in [Machine.dtns, Machine.dtn01] else target_machine
+        )
+
+        body = GlobusBodyPost(
+            source_uuid=source_name,
+            target_uuid=target_name,
+            source_dir=str(source_dir),
+            target_dir=str(target_dir),
+            label=label,
+        )
+
+        r = await self.client.post("storage/globus/transfer", data=body.model_dump())
+        new_transfer = GlobusTransferModel.model_validate(r.json())
+        transfer_id = new_transfer.transfer_id
+        r = await self.client.get(f"storage/globus/transfer/{transfer_id}")
+        json_response = r.json()
+        json_response["transfer_id"] = transfer_id
+        json_response["globus"] = self
+        transfer = AsyncGlobusTransfer.model_validate(json_response)
+        return transfer
+
+    @check_auth
+    async def transfer(self, transfer_id: str) -> GlobusTransferResult:
+        """Check on Globus transfer status
+
+        - Must select the Globus option when creating the SuperFacility key
+
+        >>> from sfapi_client import AsyncClient
+        >>> async with AsyncClient(client_id, client_secret) as client:
+        >>>     globus = client.storage.globus()
+        >>>     res = await globus.transfer(
+                        "globus-transfer-uuid"
+                )
+
+        :param str transfer_uuid: Globus UUID for the transfer
+        :return GlobusTransferResult
+        """
+        if transfer_id is None:
+            raise ValueError("Must provide a transfer_uuid")
+
+        r = await self.client.get(f"storage/globus/transfer/{transfer_id}")
+        json_response = r.json()
+        json_response["transfer_id"] = transfer_id
+        json_response["globus"] = self
+        transfer = AsyncGlobusTransfer.model_validate(json_response)
+        return transfer

src/sfapi_client/_async/users.py

@@ -1,5 +1,4 @@
-from typing import List, Optional, Callable
-from functools import wraps
+from typing import List, Optional
 
 from pydantic import ConfigDict
 
@@ -10,18 +9,7 @@
 )
 from .projects import AsyncProject, AsyncRole
 from ..exceptions import SfApiError
-
-
-def check_auth(method: Callable):
-    @wraps(method)
-    def wrapper(self, *args, **kwargs):
-        if self._client_id is None:
-            raise SfApiError(
-                f"Cannot call {self.__class__.__name__}.{method.__name__}() with an unauthenticated client."  # noqa: E501
-            )
-        return method(self, *args, **kwargs)
-
-    return wrapper
+from .._utils import check_auth
 
 
 class AsyncUser(UserBase):