wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

wandb/sdk/launch/environment/gcp_environment.py

@@ -0,0 +1,335 @@
+"""Implementation of the GCP environment for wandb launch."""
+
+import logging
+import os
+import subprocess
+from typing import Optional
+
+from wandb.sdk.launch.errors import LaunchError
+from wandb.util import get_module
+
+from ..utils import GCS_URI_RE, event_loop_thread_exec
+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__)
+
+GCP_REGION_ENV_VAR = "GOOGLE_CLOUD_REGION"
+
+
+class GcpEnvironment(AbstractEnvironment):
+    """GCP Environment.
+
+    Attributes:
+        region: The GCP region.
+    """
+
+    region: str
+
+    def __init__(
+        self,
+        region: str,
+    ) -> 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 = ""
+
+    @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)
+
+    @classmethod
+    def from_default(
+        cls,
+    ) -> "GcpEnvironment":
+        """Create a GcpEnvironment from the default configuration.
+
+        Returns:
+            GcpEnvironment: The GcpEnvironment.
+        """
+        region = get_default_region()
+        if region is None:
+            raise LaunchError(
+                "Could not create GcpEnvironment from user's gcloud configuration. "
+                "Please set the default region with `gcloud config set compute/region` "
+                "or set the environment variable {GCP_REGION_ENV_VAR}. "
+                "Alternatively, you may specify the region explicitly in your "
+                "wandb launch configuration at `$HOME/.config/wandb/launch-config.yaml`. "
+                "See https://docs.wandb.ai/guides/launch/run-agent#environments for more information."
+            )
+        return cls(region=region)
+
+    @property
+    def project(self) -> str:
+        """Get the name of the gcp project associated with the credentials.
+
+        Returns:
+            str: The name of the gcp project.
+
+        Raises:
+            LaunchError: If the launch environment cannot be verified.
+        """
+        return self._project
+
+    async 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.
+        try:
+            google_auth_default = event_loop_thread_exec(google.auth.default)
+            creds, project = await google_auth_default()
+            if not self._project:
+                self._project = project
+            _logger.debug("Refreshing GCP credentials")
+            await event_loop_thread_exec(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
+
+    async 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")
+        await self.get_credentials()
+
+    async 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)
+        cloud_storage_client = event_loop_thread_exec(google.cloud.storage.Client)
+        try:
+            credentials = await self.get_credentials()
+            storage_client = await cloud_storage_client(credentials=credentials)
+            bucket = await event_loop_thread_exec(storage_client.get_bucket)(bucket)
+        except google.api_core.exceptions.GoogleAPICallError as e:
+            raise LaunchError(
+                f"Failed verifying storage uri {uri}: bucket {bucket} does not exist."
+            ) from e
+        except google.api_core.exceptions.Forbidden as e:
+            raise LaunchError(
+                f"Failed verifying storage uri {uri}: bucket {bucket} is not accessible. Please check your permissions and try again."
+            ) from e
+
+    async 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}")
+        _err_prefix = f"Could not upload file {source} to GCS destination {destination}"
+        if not os.path.isfile(source):
+            raise LaunchError(f"{_err_prefix}: File {source} does not exist.")
+        match = GCS_URI_RE.match(destination)
+        if not match:
+            raise LaunchError(f"{_err_prefix}: Invalid GCS URI: {destination}")
+        bucket = match.group(1)
+        key = match.group(2).lstrip("/")
+        google_storage_client = event_loop_thread_exec(google.cloud.storage.Client)
+        credentials = await self.get_credentials()
+        try:
+            storage_client = await google_storage_client(credentials=credentials)
+            bucket = await event_loop_thread_exec(storage_client.bucket)(bucket)
+            blob = await event_loop_thread_exec(bucket.blob)(key)
+            await event_loop_thread_exec(blob.upload_from_filename)(source)
+        except google.api_core.exceptions.GoogleAPICallError as e:
+            resp = e.response
+            assert resp is not None
+            try:
+                message = resp.json()["error"]["message"]
+            except Exception:
+                message = str(resp)
+            raise LaunchError(f"{_err_prefix}: {message}") from e
+
+    async 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}")
+        _err_prefix = (
+            f"Could not upload directory {source} to GCS destination {destination}"
+        )
+        if not os.path.isdir(source):
+            raise LaunchError(f"{_err_prefix}: Directory {source} does not exist.")
+        match = GCS_URI_RE.match(destination)
+        if not match:
+            raise LaunchError(f"{_err_prefix}: Invalid GCS URI: {destination}")
+        bucket = match.group(1)
+        key = match.group(2).lstrip("/")
+        google_storage_client = event_loop_thread_exec(google.cloud.storage.Client)
+        credentials = await self.get_credentials()
+        try:
+            storage_client = await google_storage_client(credentials=credentials)
+            bucket = await event_loop_thread_exec(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 = await event_loop_thread_exec(bucket.blob)(gcs_path)
+                    await event_loop_thread_exec(blob.upload_from_filename)(local_path)
+        except google.api_core.exceptions.GoogleAPICallError as e:
+            resp = e.response
+            assert resp is not None
+            try:
+                message = resp.json()["error"]["message"]
+            except Exception:
+                message = str(resp)
+            raise LaunchError(f"{_err_prefix}: {message}") from e
+        except Exception as e:
+            raise LaunchError(f"{_err_prefix}: GCS upload failed: {e}") from e
+
+
+def get_gcloud_config_value(config_name: str) -> Optional[str]:
+    """Get a value from gcloud config.
+
+    Arguments:
+        config_name: The name of the config value.
+
+    Returns:
+        str: The config value, or None if the value is not set.
+    """
+    try:
+        output = subprocess.check_output(
+            ["gcloud", "config", "get-value", config_name], stderr=subprocess.STDOUT
+        )
+        value = str(output.decode("utf-8").strip())
+        if value and "unset" not in value:
+            return value
+        return None
+    except subprocess.CalledProcessError:
+        return None
+
+
+def get_default_region() -> Optional[str]:
+    """Get the default region from gcloud config or environment variables.
+
+    Returns:
+        str: The default region, or None if it cannot be determined.
+    """
+    region = get_gcloud_config_value("compute/region")
+    if not region:
+        region = os.environ.get(GCP_REGION_ENV_VAR)
+    return region

wandb/sdk/launch/environment/local_environment.py

@@ -0,0 +1,66 @@
+"""Dummy local environment implementation. This is the default environment."""
+
+from typing import Any, Dict, Union
+
+from wandb.sdk.launch.errors 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()
+
+    async def verify(self) -> None:
+        """Verify that the local environment is configured correctly."""
+        raise LaunchError("Attempted to verify LocalEnvironment.")
+
+    async 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.")
+
+    async 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.")
+
+    async 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.")
+
+    async 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/errors.py

@@ -0,0 +1,19 @@
+from wandb.errors import Error
+
+
+class LaunchError(Error):
+    """Raised when a known error occurs in wandb launch."""
+
+    pass
+
+
+class LaunchDockerError(Error):
+    """Raised when Docker daemon is not running."""
+
+    pass
+
+
+class ExecutionError(Error):
+    """Generic execution exception."""
+
+    pass

wandb/sdk/launch/git_reference.py

@@ -0,0 +1,109 @@
+"""Support for parsing GitHub URLs (which might be user provided) into constituent parts."""
+
+import re
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import Optional, Tuple, Union
+
+from wandb.sdk.launch.errors import LaunchError
+
+PREFIX_HTTPS = "https://"
+PREFIX_SSH = "git@"
+SUFFIX_GIT = ".git"
+
+
+GIT_COMMIT_REGEX = re.compile(r"[0-9a-f]{40}")
+
+
+class ReferenceType(IntEnum):
+    BRANCH = 1
+    COMMIT = 2
+
+
+def _parse_netloc(netloc: str) -> Tuple[Optional[str], Optional[str], str]:
+    """Parse netloc into username, password, and host.
+
+    github.com => None, None, "@github.com"
+    username@github.com => "username", None, "github.com"
+    username:password@github.com => "username", "password", "github.com"
+    """
+    parts = netloc.split("@", 1)
+    if len(parts) == 1:
+        return None, None, parts[0]
+    auth, host = parts
+    parts = auth.split(":", 1)
+    if len(parts) == 1:
+        return parts[0], None, host
+    return parts[0], parts[1], host
+
+
+@dataclass
+class GitReference:
+    def __init__(self, remote: str, ref: Optional[str] = None) -> None:
+        """Initialize a reference from a remote and ref.
+
+        Arguments:
+            remote: A remote URL or URI.
+            ref: A branch, tag, or commit hash.
+        """
+        self.uri = remote
+        self.ref = ref
+
+    @property
+    def url(self) -> Optional[str]:
+        return self.uri
+
+    def fetch(self, dst_dir: str) -> None:
+        """Fetch the repo into dst_dir and refine githubref based on what we learn."""
+        # We defer importing git until the last moment, because the import requires that the git
+        # executable is available on the PATH, so we only want to fail if we actually need it.
+        import git  # type: ignore
+
+        repo = git.Repo.init(dst_dir)
+        self.path = repo.working_dir
+        origin = repo.create_remote("origin", self.uri or "")
+
+        try:
+            # We fetch the origin so that we have branch and tag references
+            origin.fetch()
+        except git.exc.GitCommandError as e:
+            raise LaunchError(
+                f"Unable to fetch from git remote repository {self.url}:\n{e}"
+            )
+
+        ref: Union[git.RemoteReference, str]
+        if self.ref:
+            if self.ref in origin.refs:
+                ref = origin.refs[self.ref]
+            else:
+                ref = self.ref
+            head = repo.create_head(self.ref, ref)
+            head.checkout()
+            self.commit_hash = head.commit.hexsha
+
+        else:
+            # TODO: Is there a better way to do this?
+            default_branch = None
+            for ref in repo.references:
+                if hasattr(ref, "tag"):  # Skip tag references
+                    continue
+                refname = ref.name
+                if refname.startswith("origin/"):  # Trim off "origin/"
+                    refname = refname[7:]
+                if refname == "main":
+                    default_branch = "main"
+                    break
+                if refname == "master":
+                    default_branch = "master"
+                    # Keep looking in case we also have a main, which we let take precedence
+                    # (While the references appear to be sorted, not clear if that's guaranteed.)
+            if not default_branch:
+                raise LaunchError(
+                    f"Unable to determine branch or commit to checkout from {self.url}"
+                )
+            self.default_branch = default_branch
+            self.ref = default_branch
+            head = repo.create_head(default_branch, origin.refs[default_branch])
+            head.checkout()
+            self.commit_hash = head.commit.hexsha
+        repo.submodule_update(init=True, recursive=True)

wandb/sdk/launch/inputs/files.py

@@ -0,0 +1,148 @@
+import json
+import os
+from typing import Any, Dict
+
+import yaml
+
+from ..errors import LaunchError
+
+FILE_OVERRIDE_ENV_VAR = "WANDB_LAUNCH_FILE_OVERRIDES"
+
+
+class FileOverrides:
+    """Singleton that read file overrides json from environment variables."""
+
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = object.__new__(cls)
+            cls._instance.overrides = {}
+            cls._instance.load()
+        return cls._instance
+
+    def load(self) -> None:
+        """Load overrides from an environment variable."""
+        overrides = os.environ.get(FILE_OVERRIDE_ENV_VAR)
+        if overrides is None:
+            if f"{FILE_OVERRIDE_ENV_VAR}_0" in os.environ:
+                overrides = ""
+                idx = 0
+                while f"{FILE_OVERRIDE_ENV_VAR}_{idx}" in os.environ:
+                    overrides += os.environ[f"{FILE_OVERRIDE_ENV_VAR}_{idx}"]
+                    idx += 1
+        if overrides:
+            try:
+                contents = json.loads(overrides)
+                if not isinstance(contents, dict):
+                    raise LaunchError(f"Invalid JSON in {FILE_OVERRIDE_ENV_VAR}")
+                self.overrides = contents
+            except json.JSONDecodeError:
+                raise LaunchError(f"Invalid JSON in {FILE_OVERRIDE_ENV_VAR}")
+
+
+def config_path_is_valid(path: str) -> None:
+    """Validate a config file path.
+
+    This function checks if a given config file path is valid. A valid path
+    should meet the following criteria:
+
+    - The path must be expressed as a relative path without any upwards path
+      traversal, e.g. `../config.json`.
+    - The file specified by the path must exist.
+    - The file must have a supported extension (`.json`, `.yaml`, or `.yml`).
+
+    Args:
+        path (str): The path to validate.
+
+    Raises:
+        LaunchError: If the path is not valid.
+    """
+    if os.path.isabs(path):
+        raise LaunchError(
+            f"Invalid config path: {path}. Please provide a relative path."
+        )
+    if ".." in path:
+        raise LaunchError(
+            f"Invalid config path: {path}. Please provide a relative path "
+            "without any upward path traversal, e.g. `../config.json`."
+        )
+    path = os.path.normpath(path)
+    if not os.path.exists(path):
+        raise LaunchError(f"Invalid config path: {path}. File does not exist.")
+    if not any(path.endswith(ext) for ext in [".json", ".yaml", ".yml"]):
+        raise LaunchError(
+            f"Invalid config path: {path}. Only JSON and YAML files are supported."
+        )
+
+
+def override_file(path: str) -> None:
+    """Check for file overrides in the environment and apply them if found."""
+    file_overrides = FileOverrides()
+    if path in file_overrides.overrides:
+        overrides = file_overrides.overrides.get(path)
+        if overrides is not None:
+            config = _read_config_file(path)
+            _update_dict(config, overrides)
+            _write_config_file(path, config)
+
+
+def _write_config_file(path: str, config: Any) -> None:
+    """Write a config file to disk.
+
+    Args:
+        path (str): The path to the config file.
+        config (Any): The contents of the config file as a Python object.
+
+    Raises:
+        LaunchError: If the file extension is not supported.
+    """
+    _, ext = os.path.splitext(path)
+    if ext == ".json":
+        with open(path, "w") as f:
+            json.dump(config, f, indent=2)
+    elif ext in [".yaml", ".yml"]:
+        with open(path, "w") as f:
+            yaml.safe_dump(config, f)
+    else:
+        raise LaunchError(f"Unsupported file extension: {ext}")
+
+
+def _read_config_file(path: str) -> Any:
+    """Read a config file from disk.
+
+    Args:
+        path (str): The path to the config file.
+
+    Returns:
+        Any: The contents of the config file as a Python object.
+    """
+    _, ext = os.path.splitext(path)
+    if ext == ".json":
+        with open(
+            path,
+        ) as f:
+            return json.load(f)
+    elif ext in [".yaml", ".yml"]:
+        with open(
+            path,
+        ) as f:
+            return yaml.safe_load(f)
+    else:
+        raise LaunchError(f"Unsupported file extension: {ext}")
+
+
+def _update_dict(target: Dict, source: Dict) -> None:
+    """Update a dictionary with the contents of another dictionary.
+
+    Args:
+        target (Dict): The dictionary to update.
+        source (Dict): The dictionary to update from.
+    """
+    for key, value in source.items():
+        if isinstance(value, dict):
+            if key not in target:
+                target[key] = {}
+            _update_dict(target[key], value)
+        else:
+            target[key] = value