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

wandb/sdk/launch/builder/kaniko_builder.py

@@ -8,9 +8,12 @@ import time
 from typing import Optional
 
 import wandb
+from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
 from wandb.sdk.launch.builder.abstract import AbstractBuilder
 from wandb.sdk.launch.environment.abstract import AbstractEnvironment
+from wandb.sdk.launch.environment.azure_environment import AzureEnvironment
 from wandb.sdk.launch.registry.abstract import AbstractRegistry
+from wandb.sdk.launch.registry.azure_container_registry import AzureContainerRegistry
 from wandb.sdk.launch.registry.elastic_container_registry import (
     ElasticContainerRegistry,
 )
@@ -50,13 +53,19 @@ _DEFAULT_BUILD_TIMEOUT_SECS = 1800  # 30 minute build timeout
 
 SERVICE_ACCOUNT_NAME = os.environ.get("WANDB_LAUNCH_SERVICE_ACCOUNT_NAME", "default")
 
+if os.path.exists("/var/run/secrets/kubernetes.io/serviceaccount/namespace"):
+    with open("/var/run/secrets/kubernetes.io/serviceaccount/namespace") as f:
+        NAMESPACE = f.read().strip()
+else:
+    NAMESPACE = "wandb"
+
 
 def _wait_for_completion(
     batch_client: client.BatchV1Api, job_name: str, deadline_secs: Optional[int] = None
 ) -> bool:
     start_time = time.time()
     while True:
-        job = batch_client.read_namespaced_job_status(job_name, "wandb")
+        job = batch_client.read_namespaced_job_status(job_name, NAMESPACE)
         if job.status.succeeded is not None and job.status.succeeded >= 1:
             return True
         elif job.status.failed is not None and job.status.failed >= 1:
@@ -78,6 +87,7 @@ class KanikoBuilder(AbstractBuilder):
     build_context_store: str
     secret_name: Optional[str]
     secret_key: Optional[str]
+    image: str
 
     def __init__(
         self,
@@ -87,6 +97,7 @@ class KanikoBuilder(AbstractBuilder):
         build_context_store: str = "",
         secret_name: str = "",
         secret_key: str = "",
+        image: str = "gcr.io/kaniko-project/executor:v1.11.0",
         verify: bool = True,
     ):
         """Initialize a KanikoBuilder.
@@ -113,6 +124,7 @@ class KanikoBuilder(AbstractBuilder):
         self.build_context_store = build_context_store.rstrip("/")
         self.secret_name = secret_name
         self.secret_key = secret_key
+        self.image = image
         if verify:
             self.verify()
 
@@ -151,6 +163,7 @@ class KanikoBuilder(AbstractBuilder):
         build_job_name = config.get("build-job-name", "wandb-launch-container-build")
         secret_name = config.get("secret-name", "")
         secret_key = config.get("secret-key", "")
+        image = config.get("kaniko-image", "gcr.io/kaniko-project/executor:v1.11.0")
         return cls(
             environment,
             registry,
@@ -158,6 +171,7 @@ class KanikoBuilder(AbstractBuilder):
             build_job_name=build_job_name,
             secret_name=secret_name,
             secret_key=secret_key,
+            image=image,
             verify=verify,
         )
 
@@ -187,7 +201,7 @@ class KanikoBuilder(AbstractBuilder):
             kind="ConfigMap",
             metadata=client.V1ObjectMeta(
                 name=f"docker-config-{job_name}",
-                namespace="wandb",
+                namespace=NAMESPACE,
             ),
             data={
                 "config.json": json.dumps(
@@ -196,13 +210,13 @@ class KanikoBuilder(AbstractBuilder):
             },
             immutable=True,
         )
-        corev1_client.create_namespaced_config_map("wandb", ecr_config_map)
+        corev1_client.create_namespaced_config_map(NAMESPACE, ecr_config_map)
 
     def _delete_docker_ecr_config_map(
         self, job_name: str, client: client.CoreV1Api
     ) -> None:
         if self.secret_name:
-            client.delete_namespaced_config_map(f"docker-config-{job_name}", "wandb")
+            client.delete_namespaced_config_map(f"docker-config-{job_name}", NAMESPACE)
 
     def _upload_build_context(self, run_id: str, context_path: str) -> str:
         # creat a tar archive of the build context and upload it to s3
@@ -220,6 +234,7 @@ class KanikoBuilder(AbstractBuilder):
         self,
         launch_project: LaunchProject,
         entrypoint: EntryPoint,
+        job_tracker: Optional[JobAndRunStatusTracker] = None,
     ) -> str:
         # TODO: this should probably throw an error if the registry is a local registry
         if not self.registry:
@@ -255,35 +270,52 @@ class KanikoBuilder(AbstractBuilder):
         _, api_client = get_kube_context_and_api_client(
             kubernetes, launch_project.resource_args
         )
+        # TODO: use same client as kuberentes_runner.py
+        batch_v1 = client.BatchV1Api(api_client)
+        core_v1 = client.CoreV1Api(api_client)
+
         build_job_name = f"{self.build_job_name}-{run_id}"
 
         build_context = self._upload_build_context(run_id, context_path)
         build_job = self._create_kaniko_job(
-            build_job_name,
-            repo_uri,
-            image_uri,
-            build_context,
+            build_job_name, repo_uri, image_uri, build_context, core_v1
         )
         wandb.termlog(f"{LOG_PREFIX}Created kaniko job {build_job_name}")
 
-        # TODO: use same client as kuberentes.py
-        batch_v1 = client.BatchV1Api(api_client)
-        core_v1 = client.CoreV1Api(api_client)
-
         try:
+            if isinstance(self.registry, AzureContainerRegistry):
+                dockerfile_config_map = client.V1ConfigMap(
+                    metadata=client.V1ObjectMeta(
+                        name=f"docker-config-{build_job_name}"
+                    ),
+                    data={
+                        "config.json": json.dumps(
+                            {
+                                "credHelpers": {
+                                    f"{self.registry.registry_name}.azurecr.io": "acr-env"
+                                }
+                            }
+                        )
+                    },
+                )
+                core_v1.create_namespaced_config_map("wandb", dockerfile_config_map)
             # core_v1.create_namespaced_config_map("wandb", dockerfile_config_map)
             if self.secret_name:
                 self._create_docker_ecr_config_map(build_job_name, core_v1, repo_uri)
-            batch_v1.create_namespaced_job("wandb", build_job)
+            batch_v1.create_namespaced_job(NAMESPACE, build_job)
 
             # wait for double the job deadline since it might take time to schedule
             if not _wait_for_completion(
                 batch_v1, build_job_name, 3 * _DEFAULT_BUILD_TIMEOUT_SECS
             ):
+                if job_tracker:
+                    job_tracker.set_err_stage("build")
                 raise Exception(f"Failed to build image in kaniko for job {run_id}")
             try:
-                logs = batch_v1.read_namespaced_job_log(build_job_name, "wandb")
-                warn_failed_packages_from_build_logs(logs, image_uri)
+                logs = batch_v1.read_namespaced_job_log(build_job_name, NAMESPACE)
+                warn_failed_packages_from_build_logs(
+                    logs, image_uri, launch_project.api, job_tracker
+                )
             except Exception as e:
                 wandb.termwarn(
                     f"{LOG_PREFIX}Failed to get logs for kaniko job {build_job_name}: {e}"
@@ -298,9 +330,13 @@ class KanikoBuilder(AbstractBuilder):
             try:
                 # should we clean up the s3 build contexts? can set bucket level policy to auto deletion
                 # core_v1.delete_namespaced_config_map(config_map_name, "wandb")
+                if isinstance(self.registry, AzureContainerRegistry):
+                    core_v1.delete_namespaced_config_map(
+                        f"docker-config-{build_job_name}", "wandb"
+                    )
                 if self.secret_name:
                     self._delete_docker_ecr_config_map(build_job_name, core_v1)
-                batch_v1.delete_namespaced_job(build_job_name, "wandb")
+                batch_v1.delete_namespaced_job(build_job_name, NAMESPACE)
             except Exception as e:
                 raise LaunchError(f"Exception during Kubernetes resource clean up {e}")
 
@@ -312,6 +348,7 @@ class KanikoBuilder(AbstractBuilder):
         repository: str,
         image_tag: str,
         build_context_path: str,
+        core_client: client.CoreV1Api,
     ) -> "client.V1Job":
         env = []
         volume_mounts = []
@@ -328,6 +365,33 @@ class KanikoBuilder(AbstractBuilder):
                     value=self.registry.environment.region,
                 )
             ]
+        # TODO: Refactor all of this environment/registry
+        # specific stuff into methods of those classes.
+        if isinstance(self.environment, AzureEnvironment):
+            # Use the core api to check if the secret exists
+            try:
+                core_client.read_namespaced_secret(
+                    "azure-storage-access-key",
+                    "wandb",
+                )
+            except Exception as e:
+                raise LaunchError(
+                    "Secret azure-storage-access-key does not exist in "
+                    "namespace wandb. Please create it with the key password "
+                    "set to your azure storage access key."
+                ) from e
+            env += [
+                client.V1EnvVar(
+                    name="AZURE_STORAGE_ACCESS_KEY",
+                    value_from=client.V1EnvVarSource(
+                        secret_key_ref=client.V1SecretKeySelector(
+                            name="azure-storage-access-key",
+                            key="password",
+                        )
+                    ),
+                )
+            ]
+
         if self.secret_name and self.secret_key:
             volumes += [
                 client.V1Volume(
@@ -382,26 +446,48 @@ class KanikoBuilder(AbstractBuilder):
                     ),
                 )
             ]
-
+        if isinstance(self.registry, AzureContainerRegistry):
+            # ADd the docker config map
+            volume_mounts += [
+                client.V1VolumeMount(
+                    name="docker-config", mount_path="/kaniko/.docker/"
+                ),
+            ]
+            volumes += [
+                client.V1Volume(
+                    name="docker-config",
+                    config_map=client.V1ConfigMapVolumeSource(
+                        name=f"docker-config-{job_name}",
+                    ),
+                ),
+            ]
+        # Kaniko doesn't want https:// at the begining of the image tag.
+        destination = image_tag
+        if destination.startswith("https://"):
+            destination = destination.replace("https://", "")
         args = [
             f"--context={build_context_path}",
             "--dockerfile=Dockerfile.wandb-autogenerated",
-            f"--destination={image_tag}",
+            f"--destination={destination}",
             "--cache=true",
-            f"--cache-repo={repository}",
+            f"--cache-repo={repository.replace('https://', '')}",
             "--snapshotMode=redo",
             "--compressed-caching=false",
         ]
         container = client.V1Container(
             name="wandb-container-build",
-            image="gcr.io/kaniko-project/executor:v1.8.0",
+            image=self.image,
             args=args,
             volume_mounts=volume_mounts,
             env=env if env else None,
         )
         # Create and configure a spec section
+        labels = {"wandb": "launch"}
+        # This annotation is required to enable azure workload identity.
+        if isinstance(self.registry, AzureContainerRegistry):
+            labels["azure.workload.identity/use"] = "true"
         template = client.V1PodTemplateSpec(
-            metadata=client.V1ObjectMeta(labels={"wandb": "launch"}),
+            metadata=client.V1ObjectMeta(labels=labels),
             spec=client.V1PodSpec(
                 restart_policy="Never",
                 active_deadline_seconds=_DEFAULT_BUILD_TIMEOUT_SECS,
@@ -416,7 +502,7 @@ class KanikoBuilder(AbstractBuilder):
             api_version="batch/v1",
             kind="Job",
             metadata=client.V1ObjectMeta(
-                name=job_name, namespace="wandb", labels={"wandb": "launch"}
+                name=job_name, namespace=NAMESPACE, labels={"wandb": "launch"}
             ),
             spec=spec,
         )

wandb/sdk/launch/builder/noop.py

@@ -1,5 +1,5 @@
 """NoOp builder implementation."""
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 from wandb.sdk.launch.builder.abstract import AbstractBuilder
 from wandb.sdk.launch.environment.abstract import AbstractEnvironment
@@ -7,6 +7,7 @@ from wandb.sdk.launch.errors import LaunchError
 from wandb.sdk.launch.registry.abstract import AbstractRegistry
 
 from .._project_spec import EntryPoint, LaunchProject
+from ..agent.job_status_tracker import JobAndRunStatusTracker
 
 
 class NoOpBuilder(AbstractBuilder):
@@ -43,6 +44,7 @@ class NoOpBuilder(AbstractBuilder):
         self,
         launch_project: LaunchProject,
         entrypoint: EntryPoint,
+        job_tracker: Optional[JobAndRunStatusTracker] = None,
     ) -> str:
         """Build the image.
 

wandb/sdk/launch/environment/aws_environment.py

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

@@ -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,12 @@ 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

wandb/sdk/launch/launch.py

@@ -175,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/loader.py

@@ -54,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}"
     )
@@ -110,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}"
     )