wandb 0.13.10__py3-none-any.whl → 0.14.0__py3-none-any.whl

wandb/sdk/launch/environment/aws_environment.py

@@ -0,0 +1,276 @@
+"""Implements the AWS environment."""
+
+import logging
+import os
+import re
+from typing import Dict
+
+from wandb.sdk.launch.utils import LaunchError
+from wandb.util import get_module
+
+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__)
+
+S3_URI_RE = re.compile(r"s3://([^/]+)/(.+)")
+
+
+class AwsEnvironment(AbstractEnvironment):
+    """AWS environment."""
+
+    def __init__(
+        self,
+        region: str,
+        access_key: str,
+        secret_key: str,
+        session_token: str,
+        verify: bool = True,
+    ) -> 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
+        if verify:
+            self.verify()
+
+    @classmethod
+    def from_default(cls, region: str, verify: bool = True) -> "AwsEnvironment":
+        """Create an AWS environment from the default AWS environment.
+
+        Arguments:
+            region (str): 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()
+            region = region or session.region_name
+            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}"
+            )
+        return cls(
+            region=region,
+            access_key=access_key,
+            secret_key=secret_key,
+            session_token=session_token,
+            verify=verify,
+        )
+
+    @classmethod
+    def from_config(
+        cls, config: Dict[str, str], verify: bool = True
+    ) -> "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,
+            verify=verify,
+        )
+
+    @property
+    def region(self) -> str:
+        """The AWS region."""
+        return self._region
+
+    @region.setter
+    def region(self, region: str) -> None:
+        self._region = region
+
+    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 = self.get_session()
+            client = session.client("sts")
+            client.get_caller_identity()
+            # 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
+
+    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:
+            return boto3.Session(
+                aws_access_key_id=self._access_key,
+                aws_secret_access_key=self._secret_key,
+                aws_session_token=self._session_token,
+                region_name=self._region,
+            )
+        except botocore.exceptions.ClientError as e:
+            raise LaunchError(f"Could not create AWS session. {e}")
+
+    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}")
+        if not os.path.isfile(source):
+            raise LaunchError(f"Source {source} does not exist.")
+        match = S3_URI_RE.match(destination)
+        if not match:
+            raise LaunchError(f"Destination {destination} is not a valid s3 URI.")
+        bucket = match.group(1)
+        key = match.group(2).lstrip("/")
+        if not key:
+            key = ""
+        session = self.get_session()
+        try:
+            client = session.client("s3")
+            client.upload_file(source, bucket, key)
+        except botocore.exceptions.ClientError as e:
+            raise LaunchError(
+                f"botocore error attempting to copy {source} to {destination}. {e}"
+            ) from e
+
+    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}")
+        if not os.path.isdir(source):
+            raise LaunchError(f"Source {source} does not exist.")
+        match = S3_URI_RE.match(destination)
+        if not match:
+            raise LaunchError(f"Destination {destination} is not a valid s3 URI.")
+        bucket = match.group(1)
+        key = match.group(2).lstrip("/")
+        if not key:
+            key = ""
+        session = self.get_session()
+        try:
+            client = 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"botocore error attempting to copy {source} to {destination}. {e}"
+            ) from e
+        except Exception as e:
+            raise LaunchError(
+                f"Unexpected error attempting to copy {source} to {destination}. {e}"
+            ) from e
+
+    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"Destination {uri} is not a valid s3 URI.")
+        bucket = match.group(1)
+        try:
+            session = self.get_session()
+            client = session.client("s3")
+            client.head_bucket(Bucket=bucket)
+        except botocore.exceptions.ClientError as e:
+            raise LaunchError(
+                f"Could not verify AWS storage. Please verify that your AWS credentials are configured correctly. {e}"
+            ) from e

wandb/sdk/launch/environment/gcp_environment.py

@@ -0,0 +1,271 @@
+"""Implementation of the GCP environment for wandb launch."""
+import logging
+import os
+import re
+
+from wandb.sdk.launch.utils import LaunchError
+from wandb.util import get_module
+
+from .abstract import AbstractEnvironment
+
+google = get_module(
+    "google",
+    required="Google Cloud Platform support requires the google package. Please"
+    " install it with `pip install wandb[launch]`.",
+)
+google.cloud.compute_v1 = get_module(
+    "google.cloud.compute_v1",
+    required="Google Cloud Platform support requires the google-cloud-compute package. "
+    "Please install it with `pip install wandb[launch]`.",
+)
+google.auth.credentials = get_module(
+    "google.auth.credentials",
+    required="Google Cloud Platform support requires google-auth. "
+    "Please install it with `pip install wandb[launch]`.",
+)
+google.auth.transport.requests = get_module(
+    "google.auth.transport.requests",
+    required="Google Cloud Platform support requires google-auth. "
+    "Please install it with `pip install wandb[launch]`.",
+)
+google.api_core.exceptions = get_module(
+    "google.api_core.exceptions",
+    required="Google Cloud Platform support requires google-api-core. "
+    "Please install it with `pip install wandb[launch]`.",
+)
+google.cloud.storage = get_module(
+    "google.cloud.storage",
+    required="Google Cloud Platform support requires google-cloud-storage. "
+    "Please install it with `pip install wandb[launch].",
+)
+
+
+_logger = logging.getLogger(__name__)
+
+GCS_URI_RE = re.compile(r"gs://([^/]+)/(.+)")
+
+
+class GcpEnvironment(AbstractEnvironment):
+    """GCP Environment.
+
+    Attributes:
+        region: The GCP region.
+    """
+
+    region: str
+
+    def __init__(self, region: str, verify: bool = True) -> None:
+        """Initialize the GCP environment.
+
+        Arguments:
+            region: The GCP region.
+            verify: Whether to verify the credentials, region, and project.
+
+        Raises:
+            LaunchError: If verify is True and the environment is not properly
+                configured.
+        """
+        super().__init__()
+        _logger.info(f"Initializing GcpEnvironment in region {region}")
+        self.region: str = region
+        self._project = ""
+        if verify:
+            self.verify()
+
+    @classmethod
+    def from_config(cls, config: dict) -> "GcpEnvironment":
+        """Create a GcpEnvironment from a config dictionary.
+
+        Arguments:
+            config: The config dictionary.
+
+        Returns:
+            GcpEnvironment: The GcpEnvironment.
+        """
+        if config.get("type") != "gcp":
+            raise LaunchError(
+                f"Could not create GcpEnvironment from config. Expected type 'gcp' "
+                f"but got '{config.get('type')}'."
+            )
+        region = config.get("region", None)
+        if not region:
+            raise LaunchError(
+                "Could not create GcpEnvironment from config. Missing 'region' "
+                "field."
+            )
+        return cls(region=region)
+
+    @property
+    def project(self) -> str:
+        """Get the name of the gcp project.
+
+        The project name is determined by the credentials, so this method
+        verifies the credentials if they have not already been verified.
+
+        Returns:
+            str: The name of the gcp project.
+
+        Raises:
+            LaunchError: If the launch environment cannot be verified.
+        """
+        if not self._project:
+            raise LaunchError(
+                "This GcpEnvironment has not been verified. Please call verify() "
+                "before accessing the project."
+            )
+        return self._project
+
+    def get_credentials(self) -> google.auth.credentials.Credentials:  # type: ignore
+        """Get the GCP credentials.
+
+        Uses google.auth.default() to get the credentials. If the credentials
+        are invalid, this method will refresh them. If the credentials are
+        still invalid after refreshing, this method will raise an error.
+
+        Returns:
+            google.auth.credentials.Credentials: The GCP credentials.
+
+        Raises:
+            LaunchError: If the GCP credentials are invalid.
+        """
+        _logger.debug("Getting GCP credentials")
+        # TODO: Figure out a minimal set of scopes.
+        scopes = [
+            "https://www.googleapis.com/auth/cloud-platform",
+        ]
+        try:
+            creds, project = google.auth.default(scopes=scopes)
+            if not self._project:
+                self._project = project
+            _logger.debug("Refreshing GCP credentials")
+            creds.refresh(google.auth.transport.requests.Request())
+        except google.auth.exceptions.DefaultCredentialsError as e:
+            raise LaunchError(
+                "No Google Cloud Platform credentials found. Please run "
+                "`gcloud auth application-default login` or set the environment "
+                "variable GOOGLE_APPLICATION_CREDENTIALS to the path of a valid "
+                "service account key file."
+            ) from e
+        except google.auth.exceptions.RefreshError as e:
+            raise LaunchError(
+                "Could not refresh Google Cloud Platform credentials. Please run "
+                "`gcloud auth application-default login` or set the environment "
+                "variable GOOGLE_APPLICATION_CREDENTIALS to the path of a valid "
+                "service account key file."
+            ) from e
+        if not creds.valid:
+            raise LaunchError(
+                "Invalid Google Cloud Platform credentials. Please run "
+                "`gcloud auth application-default login` or set the environment "
+                "variable GOOGLE_APPLICATION_CREDENTIALS to the path of a valid "
+                "service account key file."
+            )
+        return creds
+
+    def verify(self) -> None:
+        """Verify the credentials, region, and project.
+
+        Credentials and region are verified by calling get_credentials(). The
+        region and is verified by calling the compute API.
+
+        Raises:
+            LaunchError: If the credentials, region, or project are invalid.
+
+        Returns:
+            None
+        """
+        _logger.debug("Verifying GCP environment")
+        creds = self.get_credentials()
+        try:
+            # Check if the region is available using the compute API.
+            compute_client = google.cloud.compute_v1.RegionsClient(credentials=creds)
+            compute_client.get(project=self.project, region=self.region)
+        except google.api_core.exceptions.NotFound as e:
+            raise LaunchError(
+                f"Region {self.region} is not available in project {self.project}."
+            ) from e
+
+    def verify_storage_uri(self, uri: str) -> None:
+        """Verify that a storage URI is valid.
+
+        Arguments:
+            uri: The storage URI.
+
+        Raises:
+            LaunchError: If the storage URI is invalid.
+        """
+        match = GCS_URI_RE.match(uri)
+        if not match:
+            raise LaunchError(f"Invalid GCS URI: {uri}")
+        bucket = match.group(1)
+        try:
+            storage_client = google.cloud.storage.Client(
+                credentials=self.get_credentials()
+            )
+            bucket = storage_client.post_bucket(bucket)
+        except google.api_core.exceptions.NotFound as e:
+            raise LaunchError(f"Bucket {bucket} does not exist.") from e
+
+    def upload_file(self, source: str, destination: str) -> None:
+        """Upload a file to GCS.
+
+        Arguments:
+            source: The path to the local file.
+            destination: The path to the GCS file.
+
+        Raises:
+            LaunchError: If the file cannot be uploaded.
+        """
+        _logger.debug(f"Uploading file {source} to {destination}")
+        if not os.path.isfile(source):
+            raise LaunchError(f"File {source} does not exist.")
+        match = GCS_URI_RE.match(destination)
+        if not match:
+            raise LaunchError(f"Invalid GCS URI: {destination}")
+        bucket = match.group(1)
+        key = match.group(2).lstrip("/")
+        try:
+            storage_client = google.cloud.storage.Client(
+                credentials=self.get_credentials()
+            )
+            bucket = storage_client.bucket(bucket)
+            blob = bucket.blob(key)
+            blob.upload_from_filename(source)
+        except google.api_core.exceptions.GoogleAPICallError as e:
+            raise LaunchError(f"Could not upload file to GCS: {e}") from e
+
+    def upload_dir(self, source: str, destination: str) -> None:
+        """Upload a directory to GCS.
+
+        Arguments:
+            source: The path to the local directory.
+            destination: The path to the GCS directory.
+
+        Raises:
+            LaunchError: If the directory cannot be uploaded.
+        """
+        _logger.debug(f"Uploading directory {source} to {destination}")
+        if not os.path.isdir(source):
+            raise LaunchError(f"Directory {source} does not exist.")
+        match = GCS_URI_RE.match(destination)
+        if not match:
+            raise LaunchError(f"Invalid GCS URI: {destination}")
+        bucket = match.group(1)
+        key = match.group(2).lstrip("/")
+        try:
+            storage_client = google.cloud.storage.Client(
+                credentials=self.get_credentials()
+            )
+            bucket = storage_client.bucket(bucket)
+            for root, _, files in os.walk(source):
+                for file in files:
+                    local_path = os.path.join(root, file)
+                    gcs_path = os.path.join(
+                        key, os.path.relpath(local_path, source)
+                    ).replace("\\", "/")
+                    blob = bucket.blob(gcs_path)
+                    blob.upload_from_filename(local_path)
+        except google.api_core.exceptions.GoogleAPICallError as e:
+            raise LaunchError(f"Could not upload directory to GCS: {e}") from e
+            raise LaunchError(f"Could not upload directory to GCS: {e}") from e
+            raise LaunchError(f"Could not upload directory to GCS: {e}") from e

wandb/sdk/launch/environment/local_environment.py

@@ -0,0 +1,65 @@
+"""Dummy local environment implementation. This is the default environment."""
+from typing import Any, Dict, Union
+
+from wandb.sdk.launch.utils import LaunchError
+
+from .abstract import AbstractEnvironment
+
+
+class LocalEnvironment(AbstractEnvironment):
+    """Local environment class."""
+
+    def __init__(self) -> None:
+        """Initialize a local environment by doing nothing."""
+        pass
+
+    @classmethod
+    def from_config(
+        cls, config: Dict[str, Union[Dict[str, Any], str]]
+    ) -> "LocalEnvironment":
+        """Create a local environment from a config.
+
+        Arguments:
+            config (dict): The config. This is ignored.
+
+        Returns:
+            LocalEnvironment: The local environment.
+        """
+        return cls()
+
+    def verify(self) -> None:
+        """Verify that the local environment is configured correctly."""
+        raise LaunchError("Attempted to verify LocalEnvironment.")
+
+    def verify_storage_uri(self, uri: str) -> None:
+        """Verify that the storage URI is configured correctly.
+
+        Arguments:
+            uri (str): The storage URI. This is ignored.
+        """
+        raise LaunchError("Attempted to verify storage uri for LocalEnvironment.")
+
+    def upload_file(self, source: str, destination: str) -> None:
+        """Upload a file from the local filesystem to storage in the environment.
+
+        Arguments:
+            source (str): The source file. This is ignored.
+            destination (str): The destination file. This is ignored.
+        """
+        raise LaunchError("Attempted to upload file for LocalEnvironment.")
+
+    def upload_dir(self, source: str, destination: str) -> None:
+        """Upload the contents of a directory from the local filesystem to the environment.
+
+        Arguments:
+            source (str): The source directory. This is ignored.
+            destination (str): The destination directory. This is ignored.
+        """
+        raise LaunchError("Attempted to upload directory for LocalEnvironment.")
+
+    def get_project(self) -> str:
+        """Get the project of the local environment.
+
+        Returns: An empty string.
+        """
+        raise LaunchError("Attempted to get project for LocalEnvironment.")

wandb/sdk/launch/github_reference.py

@@ -1,6 +1,4 @@
-"""
-Support for parsing GitHub URLs (which might be user provided) into constituent parts.
-"""
+"""Support for parsing GitHub URLs (which might be user provided) into constituent parts."""
 
 import re
 from dataclasses import dataclass
@@ -9,7 +7,7 @@ from pathlib import Path
 from typing import Optional, Tuple
 from urllib.parse import urlparse
 
-from wandb.errors import LaunchError
+from wandb.sdk.launch.utils import LaunchError
 
 PREFIX_HTTPS = "https://"
 PREFIX_SSH = "git@"
@@ -43,7 +41,6 @@ def _parse_netloc(netloc: str) -> Tuple[Optional[str], Optional[str], str]:
 
 @dataclass
 class GitHubReference:
-
     username: Optional[str] = None
     password: Optional[str] = None
     host: Optional[str] = None
@@ -107,9 +104,7 @@ class GitHubReference:
 
     @staticmethod
     def parse(uri: str) -> Optional["GitHubReference"]:
-        """
-        Attempt to parse a string as a GitHub URL.
-        """
+        """Attempt to parse a string as a GitHub URL."""
         # Special case: git@github.com:wandb/wandb.git
         ref = GitHubReference()
         if uri.startswith(PREFIX_SSH):