wandb 0.15.3__py3-none-any.whl → 0.15.5__py3-none-any.whl

wandb/sdk/launch/environment/aws_environment.py

@@ -5,7 +5,7 @@ import os
 import re
 from typing import Dict
 
-from wandb.sdk.launch.utils import LaunchError
+from wandb.sdk.launch.errors import LaunchError
 from wandb.util import get_module
 
 from .abstract import AbstractEnvironment
@@ -51,6 +51,7 @@ class AwsEnvironment(AbstractEnvironment):
         self._access_key = access_key
         self._secret_key = secret_key
         self._session_token = session_token
+        self._account = None
         if verify:
             self.verify()
 
@@ -131,7 +132,7 @@ class AwsEnvironment(AbstractEnvironment):
         try:
             session = self.get_session()
             client = session.client("sts")
-            client.get_caller_identity()
+            self._account = client.get_caller_identity().get("Account")
             # TODO: log identity details from the response
         except botocore.exceptions.ClientError as e:
             raise LaunchError(

wandb/sdk/launch/environment/azure_environment.py

@@ -0,0 +1,124 @@
+"""Implementation of AzureEnvironment class."""
+
+import re
+from typing import TYPE_CHECKING, Tuple
+
+if TYPE_CHECKING:
+    from azure.identity import DefaultAzureCredential  # type: ignore
+    from azure.storage.blob import BlobClient, BlobServiceClient  # type: ignore
+
+from wandb.util import get_module
+
+from ..errors import LaunchError
+from .abstract import AbstractEnvironment
+
+AZURE_BLOB_REGEX = re.compile(
+    r"^https://([^\.]+)\.blob\.core\.windows\.net/([^/]+)/?(.*)$"
+)
+
+
+DefaultAzureCredential = get_module(  # noqa: F811
+    "azure.identity",
+    required="The azure-identity package is required to use launch with Azure. Please install it with `pip install azure-identity`.",
+).DefaultAzureCredential
+blob = get_module(
+    "azure.storage.blob",
+    required="The azure-storage-blob package is required to use launch with Azure. Please install it with `pip install azure-storage-blob`.",
+)
+BlobClient, BlobServiceClient = blob.BlobClient, blob.BlobServiceClient  # noqa: F811
+
+
+class AzureEnvironment(AbstractEnvironment):
+    """AzureEnvironment is a helper for accessing Azure resources."""
+
+    def __init__(
+        self,
+        verify: bool = True,
+    ):
+        """Initialize an AzureEnvironment."""
+        if verify:
+            self.verify()
+
+    @classmethod
+    def from_config(cls, config: dict, verify: bool = True) -> "AzureEnvironment":
+        """Create an AzureEnvironment from a config dict."""
+        return cls(verify=verify)
+
+    @classmethod
+    def get_credentials(cls) -> DefaultAzureCredential:
+        """Get Azure credentials."""
+        try:
+            return DefaultAzureCredential()
+        except Exception as e:
+            raise LaunchError(
+                "Could not get Azure credentials. Please make sure you have "
+                "configured your Azure CLI correctly."
+            ) from e
+
+    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)
+        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)
+        except Exception as e:
+            raise LaunchError(
+                f"Could not upload file {source} to Azure blob {destination}."
+            ) from e
+
+    def upload_dir(self, source: str, destination: str) -> None:
+        """Upload a directory to Azure blob storage."""
+        raise NotImplementedError()
+
+    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
+
+    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]: The storage account and container.
+        """
+        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)

wandb/sdk/launch/environment/gcp_environment.py

@@ -3,7 +3,7 @@ import logging
 import os
 import re
 
-from wandb.sdk.launch.utils import LaunchError
+from wandb.sdk.launch.errors import LaunchError
 from wandb.util import get_module
 
 from .abstract import AbstractEnvironment
@@ -202,7 +202,7 @@ class GcpEnvironment(AbstractEnvironment):
             storage_client = google.cloud.storage.Client(
                 credentials=self.get_credentials()
             )
-            bucket = storage_client.post_bucket(bucket)
+            bucket = storage_client.get_bucket(bucket)
         except google.api_core.exceptions.NotFound as e:
             raise LaunchError(f"Bucket {bucket} does not exist.") from e
 
@@ -267,5 +267,3 @@ class GcpEnvironment(AbstractEnvironment):
                     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

@@ -1,7 +1,7 @@
 """Dummy local environment implementation. This is the default environment."""
 from typing import Any, Dict, Union
 
-from wandb.sdk.launch.utils import LaunchError
+from wandb.sdk.launch.errors import LaunchError
 
 from .abstract import AbstractEnvironment
 

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/github_reference.py

@@ -7,7 +7,7 @@ from pathlib import Path
 from typing import Optional, Tuple
 from urllib.parse import urlparse
 
-from wandb.sdk.launch.utils import LaunchError
+from wandb.sdk.launch.errors import LaunchError
 
 PREFIX_HTTPS = "https://"
 PREFIX_SSH = "git@"
@@ -58,6 +58,7 @@ class GitHubReference:
 
     ref: Optional[str] = None  # branch or commit
     ref_type: Optional[ReferenceType] = None
+    commit_hash: Optional[str] = None  # hash of commit
 
     directory: Optional[str] = None
     file: Optional[str] = None
@@ -68,6 +69,7 @@ class GitHubReference:
             self.ref_type = None
             self.ref = ref
 
+    @property
     def url_host(self) -> str:
         assert self.host
         auth = self.username or ""
@@ -77,19 +79,23 @@ class GitHubReference:
             auth += "@"
         return f"{PREFIX_HTTPS}{auth}{self.host}"
 
+    @property
     def url_organization(self) -> str:
         assert self.organization
-        return f"{self.url_host()}/{self.organization}"
+        return f"{self.url_host}/{self.organization}"
 
+    @property
     def url_repo(self) -> str:
         assert self.repo
-        return f"{self.url_organization()}/{self.repo}"
+        return f"{self.url_organization}/{self.repo}"
 
+    @property
     def repo_ssh(self) -> str:
         return f"{PREFIX_SSH}{self.host}:{self.organization}/{self.repo}{SUFFIX_GIT}"
 
+    @property
     def url(self) -> str:
-        url = self.url_repo()
+        url = self.url_repo
         if self.view:
             url += f"/{self.view}"
         if self.ref:
@@ -98,7 +104,7 @@ class GitHubReference:
                 url += f"/{self.directory}"
             if self.file:
                 url += f"/{self.file}"
-        elif self.path:
+        if self.path:
             url += f"/{self.path}"
         return url
 
@@ -127,18 +133,21 @@ class GitHubReference:
         ref.username, ref.password, ref.host = _parse_netloc(parsed.netloc)
 
         parts = parsed.path.split("/")
-        if len(parts) > 1:
-            if parts[1] == "orgs" and len(parts) > 2:
-                ref.organization = parts[2]
-            else:
-                ref.organization = parts[1]
-                if len(parts) > 2:
-                    repo = parts[2]
-                    if repo.endswith(SUFFIX_GIT):
-                        repo = repo[: -len(SUFFIX_GIT)]
-                    ref.repo = repo
-                    ref.view = parts[3] if len(parts) > 3 else None
-                    ref.path = "/".join(parts[4:])
+        if len(parts) < 2:
+            return ref
+        if parts[1] == "orgs" and len(parts) > 2:
+            ref.organization = parts[2]
+            return ref
+        ref.organization = parts[1]
+        if len(parts) < 3:
+            return ref
+        repo = parts[2]
+        if repo.endswith(SUFFIX_GIT):
+            repo = repo[: -len(SUFFIX_GIT)]
+        ref.repo = repo
+        ref.view = parts[3] if len(parts) > 3 else None
+        ref.path = "/".join(parts[4:])
+
         return ref
 
     def fetch(self, dst_dir: str) -> None:
@@ -148,7 +157,7 @@ class GitHubReference:
         import git  # type: ignore
 
         repo = git.Repo.init(dst_dir)
-        origin = repo.create_remote("origin", self.url_repo())
+        origin = repo.create_remote("origin", self.url_repo)
 
         # We fetch the origin so that we have branch and tag references
         origin.fetch(depth=1)
@@ -165,6 +174,7 @@ class GitHubReference:
                     self.path = self.path[len(first_segment) + 1 :]
                 head = repo.create_head(first_segment, commit)
                 head.checkout()
+                self.commit_hash = head.commit.hexsha
             except ValueError:
                 # Apparently it just looked like a commit
                 pass
@@ -188,6 +198,7 @@ class GitHubReference:
                         self.path = self.path[len(refname) + 1 :]
                     head = repo.create_head(branch, origin.refs[branch])
                     head.checkout()
+                    self.commit_hash = head.commit.hexsha
                     break
 
         # Must be on default branch. Try to figure out what that is.
@@ -209,11 +220,13 @@ class GitHubReference:
                     # (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()}"
+                    f"Unable to determine branch or commit to checkout from {self.url}"
                 )
             self.default_branch = 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)
 
         # Now that we've checked something out, try to extract directory and file from what remains
         self._update_path(dst_dir)

wandb/sdk/launch/launch.py

@@ -11,13 +11,12 @@ from . import loader
 from ._project_spec import create_project_from_spec, fetch_and_validate_project
 from .agent import LaunchAgent
 from .builder.build import construct_builder_args
+from .errors import ExecutionError, LaunchError
 from .runner.abstract import AbstractRun
 from .utils import (
     LAUNCH_CONFIG_FILE,
     LAUNCH_DEFAULT_PROJECT,
     PROJECT_SYNCHRONOUS,
-    ExecutionError,
-    LaunchError,
     construct_launch_spec,
     validate_launch_spec_source,
 )
@@ -76,14 +75,10 @@ def resolve_agent_config(  # noqa: C901
         user_set_project = True
     if os.environ.get("WANDB_ENTITY") is not None:
         resolved_config.update({"entity": os.environ.get("WANDB_ENTITY")})
-    if os.environ.get("WANDB_API_KEY") is not None:
-        resolved_config.update({"api_key": os.environ.get("WANDB_API_KEY")})
     if os.environ.get("WANDB_LAUNCH_MAX_JOBS") is not None:
         resolved_config.update(
             {"max_jobs": int(os.environ.get("WANDB_LAUNCH_MAX_JOBS", 1))}
         )
-    if os.environ.get("WANDB_BASE_URL") is not None:
-        resolved_config.update({"base_url": os.environ.get("WANDB_BASE_URL")})
 
     if project is not None:
         resolved_config.update({"project": project})
@@ -104,7 +99,7 @@ def resolve_agent_config(  # noqa: C901
                 + " (expected str). Specify multiple queues with the 'queues' key"
             )
 
-    keys = ["api_key", "base_url", "project", "entity"]
+    keys = ["project", "entity"]
     settings = {
         k: resolved_config.get(k) for k in keys if resolved_config.get(k) is not None
     }
@@ -180,7 +175,7 @@ def _run(
     builder = loader.builder_from_config(build_config, environment, registry)
     backend = loader.runner_from_config(resource, api, runner_config, environment)
     if backend:
-        submitted_run = backend.run(launch_project, builder)
+        submitted_run = backend.run(launch_project, builder, None)
         # this check will always pass, run is only optional in the agent case where
         # a run queue id is present on the backend config
         assert submitted_run

wandb/sdk/launch/launch_add.py

@@ -6,10 +6,10 @@ import wandb.apis.public as public
 from wandb.apis.internal import Api
 from wandb.sdk.launch._project_spec import create_project_from_spec
 from wandb.sdk.launch.builder.build import build_image_from_project
+from wandb.sdk.launch.errors import LaunchError
 from wandb.sdk.launch.utils import (
     LAUNCH_DEFAULT_PROJECT,
     LOG_PREFIX,
-    LaunchError,
     construct_launch_spec,
     validate_launch_spec_source,
 )
@@ -43,6 +43,7 @@ def launch_add(
     run_id: Optional[str] = None,
     build: Optional[bool] = False,
     repository: Optional[str] = None,
+    sweep_id: Optional[str] = None,
     author: Optional[str] = None,
 ) -> "public.QueuedRun":
     """Enqueue a W&B launch experiment. With either a source uri, job or docker_image.
@@ -111,6 +112,7 @@ def launch_add(
         run_id=run_id,
         build=build,
         repository=repository,
+        sweep_id=sweep_id,
         author=author,
     )
 
@@ -133,6 +135,7 @@ def _launch_add(
     run_id: Optional[str] = None,
     build: Optional[bool] = False,
     repository: Optional[str] = None,
+    sweep_id: Optional[str] = None,
     author: Optional[str] = None,
 ) -> "public.QueuedRun":
     launch_spec = construct_launch_spec(
@@ -151,6 +154,7 @@ def _launch_add(
         run_id,
         repository,
         author,
+        sweep_id,
     )
 
     if build:
@@ -207,7 +211,7 @@ def _launch_add(
     container_job = False
     if job:
         job_artifact = public_api.job(job)
-        if job_artifact._source_info.get("source_type") == "image":
+        if job_artifact._job_info.get("source_type") == "image":
             container_job = True
 
     queued_run = public_api.queued_run(

wandb/sdk/launch/loader.py

@@ -2,7 +2,7 @@
 from typing import Any, Dict, Optional
 
 from wandb.apis.internal import Api
-from wandb.sdk.launch.utils import LaunchError
+from wandb.sdk.launch.errors import LaunchError
 
 from .builder.abstract import AbstractBuilder
 from .environment.abstract import AbstractEnvironment
@@ -42,6 +42,10 @@ def environment_from_config(config: Optional[Dict[str, Any]]) -> AbstractEnviron
         raise LaunchError(
             "Could not create environment from config. Environment type not specified!"
         )
+    if env_type == "local":
+        from .environment.local_environment import LocalEnvironment
+
+        return LocalEnvironment.from_config(config)
     if env_type == "aws":
         from .environment.aws_environment import AwsEnvironment
 
@@ -50,6 +54,10 @@ def environment_from_config(config: Optional[Dict[str, Any]]) -> AbstractEnviron
         from .environment.gcp_environment import GcpEnvironment
 
         return GcpEnvironment.from_config(config)
+    if env_type == "azure":
+        from .environment.azure_environment import AzureEnvironment
+
+        return AzureEnvironment.from_config(config)
     raise LaunchError(
         f"Could not create environment from config. Invalid type: {env_type}"
     )
@@ -80,7 +88,7 @@ def registry_from_config(
 
         return LocalRegistry()  # This is the default, dummy registry.
     registry_type = config.get("type")
-    if registry_type is None:
+    if registry_type is None or registry_type == "local":
         from .registry.local_registry import LocalRegistry
 
         return LocalRegistry()  # This is the default, dummy registry.
@@ -106,6 +114,17 @@ def registry_from_config(
         from .registry.google_artifact_registry import GoogleArtifactRegistry
 
         return GoogleArtifactRegistry.from_config(config, environment)
+    if registry_type == "acr":
+        from .environment.azure_environment import AzureEnvironment
+
+        if not isinstance(environment, AzureEnvironment):
+            raise LaunchError(
+                "Could not create ACR registry. "
+                "Environment must be an instance of AzureEnvironment."
+            )
+        from .registry.azure_container_registry import AzureContainerRegistry
+
+        return AzureContainerRegistry.from_config(config, environment)
     raise LaunchError(
         f"Could not create registry from config. Invalid registry type: {registry_type}"
     )

wandb/sdk/launch/registry/azure_container_registry.py

@@ -0,0 +1,132 @@
+"""Implementation of AzureContainerRegistry class."""
+import re
+from typing import TYPE_CHECKING, Tuple
+
+from wandb.util import get_module
+
+from ..environment.abstract import AbstractEnvironment
+from ..environment.azure_environment import AzureEnvironment
+from ..errors import LaunchError
+from .abstract import AbstractRegistry
+
+if TYPE_CHECKING:
+    from azure.containerregistry import ContainerRegistryClient  # type: ignore
+    from azure.core.exceptions import ResourceNotFoundError  # type: ignore
+
+
+ContainerRegistryClient = get_module(  # noqa: F811
+    "azure.containerregistry",
+    required="The azure-containerregistry package is required to use launch with Azure. Please install it with `pip install azure-containerregistry`.",
+).ContainerRegistryClient
+
+ResourceNotFoundError = get_module(  # noqa: F811
+    "azure.core.exceptions",
+    required="The azure-core package is required to use launch with Azure. Please install it with `pip install azure-core`.",
+).ResourceNotFoundError
+
+
+class AzureContainerRegistry(AbstractRegistry):
+    """Helper for accessing Azure Container Registry resources."""
+
+    def __init__(
+        self,
+        environment: AzureEnvironment,
+        uri: str,
+        verify: bool = True,
+    ):
+        """Initialize an AzureContainerRegistry."""
+        self.environment = environment
+        self.uri = uri
+        if verify:
+            self.verify()
+
+    @classmethod
+    def from_config(
+        cls, config: dict, environment: AbstractEnvironment, verify: bool = True
+    ) -> "AzureContainerRegistry":
+        """Create an AzureContainerRegistry from a config dict.
+
+        Args:
+            config (dict): The config dict.
+            environment (AbstractEnvironment): The environment to use.
+            verify (bool, optional): Whether to verify the registry. Defaults to True.
+
+        Returns:
+            AzureContainerRegistry: The registry.
+
+        Raises:
+            LaunchError: If the config is invalid.
+        """
+        if not isinstance(environment, AzureEnvironment):
+            raise LaunchError(
+                "AzureContainerRegistry requires an AzureEnvironment to be passed in."
+            )
+        uri = config.get("uri")
+        if uri is None:
+            raise LaunchError(
+                "Please specify a registry name to use under the registry.uri."
+            )
+        return cls(
+            uri=uri,
+            environment=environment,
+            verify=verify,
+        )
+
+    def get_username_password(self) -> Tuple[str, str]:
+        """Get username and password for container registry."""
+        raise NotImplementedError
+
+    def check_image_exists(self, image_uri: str) -> bool:
+        """Check if image exists in container registry.
+
+        Args:
+            image_uri (str): Image URI to check.
+
+        Returns:
+            bool: True if image exists, False otherwise.
+        """
+        credential = self.environment.get_credentials()
+        registry, repository, tag = self.parse_azurecr_uri(image_uri)
+        client = ContainerRegistryClient(f"https://{registry}.azurecr.io", credential)
+        try:
+            client.get_manifest_properties(repository, tag)
+            return True
+        except ResourceNotFoundError:
+            return False
+        except Exception as e:
+            raise LaunchError(
+                f"Unable to check if image exists in Azure Container Registry: {e}"
+            ) from e
+
+    def get_repo_uri(self) -> str:
+        return self.uri
+
+    def verify(self) -> None:
+        try:
+            _ = self.registry_name
+        except Exception as e:
+            raise LaunchError(f"Unable to verify Azure Container Registry: {e}") from e
+
+    @property
+    def registry_name(self) -> str:
+        """Get registry name."""
+        return self.parse_azurecr_uri(self.uri)[0]
+
+    @staticmethod
+    def parse_azurecr_uri(uri: str) -> Tuple[str, str, str]:
+        """Parse an Azure Container Registry URI.
+
+        Args:
+            uri (str): URI to parse.
+
+        Returns:
+            Tuple[str, str, str]: Tuple of registry name, repository name, and tag.
+
+        Raises:
+            LaunchError: If unable to parse URI.
+        """
+        regex = r"(?:https://)([\w]+)\.azurecr\.io/([\w\-]+):?(.*)"
+        match = re.match(regex, uri)
+        if match is None:
+            raise LaunchError(f"Unable to parse Azure Container Registry URI: {uri}")
+        return match.group(1), match.group(2), match.group(3)