wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

wandb/sdk/launch/environment/aws_environment.py

@@ -0,0 +1,322 @@
+"""Implements the AWS environment."""
+
+import logging
+import os
+from typing import Dict, Optional
+
+from wandb.sdk.launch.errors import LaunchError
+from wandb.util import get_module
+
+from ..utils import ARN_PARTITION_RE, S3_URI_RE, event_loop_thread_exec
+from .abstract import AbstractEnvironment
+
+boto3 = get_module(
+    "boto3",
+    required="AWS environment requires boto3 to be installed. Please install "
+    "it with `pip install wandb[launch]`.",
+)
+botocore = get_module(
+    "botocore",
+    required="AWS environment requires botocore to be installed. Please install "
+    "it with `pip install wandb[launch]`.",
+)
+
+_logger = logging.getLogger(__name__)
+
+
+class AwsEnvironment(AbstractEnvironment):
+    """AWS environment."""
+
+    def __init__(
+        self,
+        region: str,
+        access_key: str,
+        secret_key: str,
+        session_token: str,
+    ) -> None:
+        """Initialize the AWS environment.
+
+        Arguments:
+            region (str): The AWS region.
+
+        Raises:
+            LaunchError: If the AWS environment is not configured correctly.
+        """
+        super().__init__()
+        _logger.info(f"Initializing AWS environment in region {region}.")
+        self._region = region
+        self._access_key = access_key
+        self._secret_key = secret_key
+        self._session_token = session_token
+        self._account = None
+        self._partition = None
+
+    @classmethod
+    def from_default(cls, region: Optional[str] = None) -> "AwsEnvironment":
+        """Create an AWS environment from the default AWS environment.
+
+        Arguments:
+            region (str, optional): The AWS region.
+            verify (bool, optional): Whether to verify the AWS environment. Defaults to True.
+
+        Returns:
+            AwsEnvironment: The AWS environment.
+        """
+        _logger.info("Creating AWS environment from default credentials.")
+        try:
+            session = boto3.Session()
+            if hasattr(session, "region"):
+                region = region or session.region
+            region = region or os.environ.get("AWS_REGION")
+            credentials = session.get_credentials()
+            if not credentials:
+                raise LaunchError(
+                    "Could not create AWS environment from default environment. Please verify that your AWS credentials are configured correctly."
+                )
+            access_key = credentials.access_key
+            secret_key = credentials.secret_key
+            session_token = credentials.token
+        except botocore.client.ClientError as e:
+            raise LaunchError(
+                f"Could not create AWS environment from default environment. Please verify that your AWS credentials are configured correctly. {e}"
+            )
+        if not region:
+            raise LaunchError(
+                "Could not create AWS environment from default environment. Region not specified."
+            )
+        return cls(
+            region=region,
+            access_key=access_key,
+            secret_key=secret_key,
+            session_token=session_token,
+        )
+
+    @classmethod
+    def from_config(
+        cls,
+        config: Dict[str, str],
+    ) -> "AwsEnvironment":
+        """Create an AWS environment from the default AWS environment.
+
+        Arguments:
+            config (dict): Configuration dictionary.
+            verify (bool, optional): Whether to verify the AWS environment. Defaults to True.
+
+        Returns:
+            AwsEnvironment: The AWS environment.
+        """
+        region = str(config.get("region", ""))
+        if not region:
+            raise LaunchError(
+                "Could not create AWS environment from config. Region not specified."
+            )
+        return cls.from_default(
+            region=region,
+        )
+
+    @property
+    def region(self) -> str:
+        """The AWS region."""
+        return self._region
+
+    @region.setter
+    def region(self, region: str) -> None:
+        self._region = region
+
+    async def get_partition(self) -> str:
+        """Set the partition for the AWS environment."""
+        try:
+            session = await self.get_session()
+            client = await event_loop_thread_exec(session.client)("sts")
+            get_caller_identity = event_loop_thread_exec(client.get_caller_identity)
+            identity = await get_caller_identity()
+            arn = identity.get("Arn")
+            if not arn:
+                raise LaunchError(
+                    "Could not set partition for AWS environment. ARN not found."
+                )
+            matched_partition = ARN_PARTITION_RE.match(arn)
+            if not matched_partition:
+                raise LaunchError(
+                    f"Could not set partition for AWS environment. ARN {arn} is not valid."
+                )
+            partition = matched_partition.group(1)
+            return partition
+        except botocore.exceptions.ClientError as e:
+            raise LaunchError(
+                f"Could not set partition for AWS environment. {e}"
+            ) from e
+
+    async def verify(self) -> None:
+        """Verify that the AWS environment is configured correctly.
+
+        Raises:
+            LaunchError: If the AWS environment is not configured correctly.
+        """
+        _logger.debug("Verifying AWS environment.")
+        try:
+            session = await self.get_session()
+            client = await event_loop_thread_exec(session.client)("sts")
+            get_caller_identity = event_loop_thread_exec(client.get_caller_identity)
+            self._account = (await get_caller_identity()).get("Account")
+            # TODO: log identity details from the response
+        except botocore.exceptions.ClientError as e:
+            raise LaunchError(
+                f"Could not verify AWS environment. Please verify that your AWS credentials are configured correctly. {e}"
+            ) from e
+
+    async def get_session(self) -> "boto3.Session":  # type: ignore
+        """Get an AWS session.
+
+        Returns:
+            boto3.Session: The AWS session.
+
+        Raises:
+            LaunchError: If the AWS session could not be created.
+        """
+        _logger.debug(f"Creating AWS session in region {self._region}")
+        try:
+            session = event_loop_thread_exec(boto3.Session)
+            return await session(
+                region_name=self._region,
+                aws_access_key_id=self._access_key,
+                aws_secret_access_key=self._secret_key,
+                aws_session_token=self._session_token,
+            )
+        except botocore.exceptions.ClientError as e:
+            raise LaunchError(f"Could not create AWS session. {e}")
+
+    async def upload_file(self, source: str, destination: str) -> None:
+        """Upload a file to s3 from local storage.
+
+        The destination is a valid s3 URI, e.g. s3://bucket/key and will
+        be used as a prefix for the uploaded file.  Only the filename of the source
+        is kept in the upload key.  So if the source is "foo/bar" and the
+        destination is "s3://bucket/key", the file "foo/bar" will be uploaded
+        to "s3://bucket/key/bar".
+
+        Arguments:
+            source (str): The path to the file or directory.
+            destination (str): The uri of the storage destination. This should
+                be a valid s3 URI, e.g. s3://bucket/key.
+
+        Raises:
+            LaunchError: If the copy fails, the source path does not exist, or the
+                destination is not a valid s3 URI, or the upload fails.
+        """
+        _logger.debug(f"Uploading {source} to {destination}")
+        _err_prefix = f"Error attempting to copy {source} to {destination}."
+        if not os.path.isfile(source):
+            raise LaunchError(f"{_err_prefix}: Source {source} does not exist.")
+        match = S3_URI_RE.match(destination)
+        if not match:
+            raise LaunchError(
+                f"{_err_prefix}: Destination {destination} is not a valid s3 URI."
+            )
+        bucket = match.group(1)
+        key = match.group(2).lstrip("/")
+        if not key:
+            key = ""
+        session = await self.get_session()
+        try:
+            client = await event_loop_thread_exec(session.client)("s3")
+            client.upload_file(source, bucket, key)
+        except botocore.exceptions.ClientError as e:
+            raise LaunchError(
+                f"{_err_prefix}: botocore error attempting to copy {source} to {destination}. {e}"
+            )
+
+    async def upload_dir(self, source: str, destination: str) -> None:
+        """Upload a directory to s3 from local storage.
+
+        The upload will place the contents of the source directory in the destination
+        with the same directory structure. So if the source is "foo/bar" and the
+        destination is "s3://bucket/key", the contents of "foo/bar" will be uploaded
+        to "s3://bucket/key/bar".
+
+        Arguments:
+            source (str): The path to the file or directory.
+            destination (str): The URI of the storage.
+            recursive (bool, optional): If True, copy the directory recursively. Defaults to False.
+
+        Raises:
+            LaunchError: If the copy fails, the source path does not exist, or the
+                destination is not a valid s3 URI.
+        """
+        _logger.debug(f"Uploading {source} to {destination}")
+        _err_prefix = f"Error attempting to copy {source} to {destination}."
+        if not os.path.isdir(source):
+            raise LaunchError(f"{_err_prefix}: Source {source} does not exist.")
+        match = S3_URI_RE.match(destination)
+        if not match:
+            raise LaunchError(
+                f"{_err_prefix}: Destination {destination} is not a valid s3 URI."
+            )
+        bucket = match.group(1)
+        key = match.group(2).lstrip("/")
+        if not key:
+            key = ""
+        session = await self.get_session()
+        try:
+            client = await event_loop_thread_exec(session.client)("s3")
+            for path, _, files in os.walk(source):
+                for file in files:
+                    abs_path = os.path.join(path, file)
+                    key_path = (
+                        abs_path.replace(source, "").replace("\\", "/").lstrip("/")
+                    )
+                    client.upload_file(
+                        abs_path,
+                        bucket,
+                        key_path,
+                    )
+        except botocore.exceptions.ClientError as e:
+            raise LaunchError(
+                f"{_err_prefix}: botocore error attempting to copy {source} to {destination}. {e}"
+            ) from e
+        except Exception as e:
+            raise LaunchError(
+                f"{_err_prefix}: Unexpected error attempting to copy {source} to {destination}. {e}"
+            ) from e
+
+    async def verify_storage_uri(self, uri: str) -> None:
+        """Verify that s3 storage is configured correctly.
+
+        This will check that the bucket exists and that the credentials are
+        configured correctly.
+
+        Arguments:
+            uri (str): The URI of the storage.
+
+        Raises:
+            LaunchError: If the storage is not configured correctly or the URI is
+                not a valid s3 URI.
+
+        Returns:
+            None
+        """
+        _logger.debug(f"Verifying storage {uri}")
+        match = S3_URI_RE.match(uri)
+        if not match:
+            raise LaunchError(
+                f"Failed to validate storage uri: {uri} is not a valid s3 URI."
+            )
+        bucket = match.group(1)
+        try:
+            session = await self.get_session()
+            client = await event_loop_thread_exec(session.client)("s3")
+            client.head_bucket(Bucket=bucket)
+        except botocore.exceptions.ClientError as e:
+            if e.response["Error"]["Code"] == "404":
+                raise LaunchError(
+                    f"Could not verify AWS storage uri {uri}. Bucket {bucket} does not exist."
+                )
+            if e.response["Error"]["Code"] == "403":
+                raise LaunchError(
+                    f"Could not verify AWS storage uri {uri}. "
+                    "Bucket {bucket} is not accessible. Please check that this "
+                    "client is authenticated with permission to access the bucket."
+                )
+            raise LaunchError(
+                f"Failed to verify AWS storage uri {uri}. Response: {e.response} Please verify that your AWS credentials are configured correctly."
+            )

wandb/sdk/launch/environment/azure_environment.py

@@ -0,0 +1,105 @@
+"""Implementation of AzureEnvironment class."""
+
+from typing import Tuple
+
+from azure.core.exceptions import HttpResponseError  # type: ignore
+from azure.identity import DefaultAzureCredential  # type: ignore
+from azure.storage.blob import BlobClient, BlobServiceClient  # type: ignore
+
+from ..errors import LaunchError
+from ..utils import AZURE_BLOB_REGEX
+from .abstract import AbstractEnvironment
+
+
+class AzureEnvironment(AbstractEnvironment):
+    """AzureEnvironment is a helper for accessing Azure resources."""
+
+    def __init__(
+        self,
+    ) -> None:
+        """Initialize an AzureEnvironment."""
+
+    @classmethod
+    def from_config(cls, config: dict, verify: bool = True) -> "AzureEnvironment":
+        """Create an AzureEnvironment from a config dict."""
+        return cls()
+
+    @classmethod
+    def get_credentials(cls) -> DefaultAzureCredential:
+        """Get Azure credentials."""
+        try:
+            return DefaultAzureCredential()
+        except Exception as e:
+            raise LaunchError(
+                f"Could not get Azure credentials. Please make sure you have "
+                f"configured your Azure CLI correctly.\n{e}"
+            ) from e
+
+    async def upload_file(self, source: str, destination: str) -> None:
+        """Upload a file to Azure blob storage.
+
+        Arguments:
+            source (str): The path to the file to upload.
+            destination (str): The destination path in Azure blob storage. Ex:
+                https://<storage_account>.blob.core.windows.net/<storage_container>/<path>
+        Raise:
+            LaunchError: If the file could not be uploaded.
+        """
+        storage_account, storage_container, path = self.parse_uri(destination)
+        _err_prefix = f"Could not upload file {source} to Azure blob {destination}"
+        creds = self.get_credentials()
+        try:
+            client = BlobClient(
+                f"https://{storage_account}.blob.core.windows.net",
+                storage_container,
+                path,
+                credential=creds,
+            )
+            with open(source, "rb") as f:
+                client.upload_blob(f, overwrite=True)
+        except HttpResponseError as e:
+            raise LaunchError(f"{_err_prefix}: {e.message}") from e
+        except Exception as e:
+            raise LaunchError(f"{_err_prefix}: {e.__class__.__name__}: {e}") from e
+
+    async def upload_dir(self, source: str, destination: str) -> None:
+        """Upload a directory to Azure blob storage."""
+        raise NotImplementedError()
+
+    async def verify_storage_uri(self, uri: str) -> None:
+        """Verify that the given blob storage prefix exists.
+
+        Args:
+            uri (str): The URI to verify.
+        """
+        creds = self.get_credentials()
+        storage_account, storage_container, _ = self.parse_uri(uri)
+        try:
+            client = BlobServiceClient(
+                f"https://{storage_account}.blob.core.windows.net",
+                credential=creds,
+            )
+            client.get_container_client(storage_container)
+        except Exception as e:
+            raise LaunchError(
+                f"Could not verify storage URI {uri} in container {storage_container}."
+            ) from e
+
+    async def verify(self) -> None:
+        """Verify that the AzureEnvironment is valid."""
+        self.get_credentials()
+
+    @staticmethod
+    def parse_uri(uri: str) -> Tuple[str, str, str]:
+        """Parse an Azure blob storage URI into a storage account and container.
+
+        Args:
+            uri (str): The URI to parse.
+
+        Returns:
+            Tuple[str, str, prefix]: The storage account, container, and path.
+        """
+        match = AZURE_BLOB_REGEX.match(uri)
+        if match is None:
+            raise LaunchError(f"Could not parse Azure blob URI {uri}.")
+        return match.group(1), match.group(2), match.group(3)