wandb 0.16.3__py3-none-any.whl → 0.16.5__py3-none-any.whl

wandb/sdk/launch/builder/docker_builder.py

@@ -11,18 +11,13 @@ from wandb.sdk.launch.builder.build import registry_from_uri
 from wandb.sdk.launch.environment.abstract import AbstractEnvironment
 from wandb.sdk.launch.registry.abstract import AbstractRegistry
 
-from .._project_spec import (
-    EntryPoint,
-    LaunchProject,
-    create_metadata_file,
-    get_entry_point_command,
-)
+from .._project_spec import EntryPoint, LaunchProject
 from ..errors import LaunchDockerError, LaunchError
+from ..registry.anon import AnonynmousRegistry
 from ..registry.local_registry import LocalRegistry
 from ..utils import (
     LOG_PREFIX,
     event_loop_thread_exec,
-    sanitize_wandb_api_key,
     warn_failed_packages_from_build_logs,
 )
 from .build import (
@@ -106,6 +101,8 @@ class DockerBuilder(AbstractBuilder):
         """Login to the registry."""
         if isinstance(self.registry, LocalRegistry):
             _logger.info(f"{LOG_PREFIX}No registry configured, skipping login.")
+        elif isinstance(self.registry, AnonynmousRegistry):
+            _logger.info(f"{LOG_PREFIX}Anonymous registry, skipping login.")
         else:
             username, password = await self.registry.get_username_password()
             login = event_loop_thread_exec(docker.login)
@@ -155,14 +152,6 @@ class DockerBuilder(AbstractBuilder):
             f"image {image_uri} does not already exist in repository, building."
         )
 
-        entry_cmd = get_entry_point_command(entrypoint, launch_project.override_args)
-
-        create_metadata_file(
-            launch_project,
-            image_uri,
-            sanitize_wandb_api_key(" ".join(entry_cmd)),
-            dockerfile_str,
-        )
         build_ctx_path = _create_docker_build_ctx(launch_project, dockerfile_str)
         dockerfile = os.path.join(build_ctx_path, _WANDB_DOCKERFILE_NAME)
         try:

wandb/sdk/launch/builder/kaniko_builder.py

@@ -3,6 +3,7 @@ import base64
 import json
 import logging
 import os
+import shutil
 import tarfile
 import tempfile
 import time
@@ -23,17 +24,11 @@ from wandb.sdk.launch.registry.elastic_container_registry import (
 from wandb.sdk.launch.registry.google_artifact_registry import GoogleArtifactRegistry
 from wandb.util import get_module
 
-from .._project_spec import (
-    EntryPoint,
-    LaunchProject,
-    create_metadata_file,
-    get_entry_point_command,
-)
+from .._project_spec import EntryPoint, LaunchProject
 from ..errors import LaunchError
 from ..utils import (
     LOG_PREFIX,
     get_kube_context_and_api_client,
-    sanitize_wandb_api_key,
     warn_failed_packages_from_build_logs,
 )
 from .build import (
@@ -56,6 +51,14 @@ _logger = logging.getLogger(__name__)
 _DEFAULT_BUILD_TIMEOUT_SECS = 1800  # 30 minute build timeout
 
 SERVICE_ACCOUNT_NAME = os.environ.get("WANDB_LAUNCH_SERVICE_ACCOUNT_NAME", "default")
+PVC_NAME = os.environ.get("WANDB_LAUNCH_KANIKO_PVC_NAME")
+PVC_MOUNT_PATH = (
+    os.environ.get("WANDB_LAUNCH_KANIKO_PVC_MOUNT_PATH", "/kaniko").rstrip("/")
+    if PVC_NAME
+    else None
+)
+DOCKER_CONFIG_SECRET = os.environ.get("WANDB_LAUNCH_KANIKO_AUTH_SECRET")
+
 
 if os.path.exists("/var/run/secrets/kubernetes.io/serviceaccount/namespace"):
     with open("/var/run/secrets/kubernetes.io/serviceaccount/namespace") as f:
@@ -115,12 +118,6 @@ class KanikoBuilder(AbstractBuilder):
             verify (bool, optional): Whether to verify the functionality of the builder.
                 Defaults to True.
         """
-        if build_context_store is None:
-            raise LaunchError(
-                "You are required to specify an external build "
-                "context store for Kaniko builds. Please specify a  storage url "
-                "in the 'build-context-store' field of your builder config."
-            )
         self.environment = environment
         self.registry = registry
         self.build_job_name = build_job_name
@@ -154,13 +151,16 @@ class KanikoBuilder(AbstractBuilder):
             raise LaunchError(
                 "Builder config must include 'type':'kaniko' to create a KanikoBuilder."
             )
-        build_context_store = config.get("build-context-store")
+        build_context_store = config.get("build-context-store", "")
         if build_context_store is None:
-            raise LaunchError(
-                "You are required to specify an external build "
-                "context store for Kaniko builds. Please specify a  "
-                "storage url in the 'build_context_store' field of your builder config."
-            )
+            if not PVC_MOUNT_PATH:
+                raise LaunchError(
+                    "You must specify a build context store for kaniko builds. "
+                    "You can set builder.build-context-store in your agent config "
+                    "to a valid s3, gcs, or azure blog storage URI. Or, configure "
+                    "a persistent volume claim through the agent helm chart: "
+                    "https://github.com/wandb/helm-charts/tree/main/charts/launch-agent"
+                )
         build_job_name = config.get("build-job-name", "wandb-launch-container-build")
         secret_name = config.get("secret-name", "")
         secret_key = config.get("secret-key", "")
@@ -170,6 +170,7 @@ class KanikoBuilder(AbstractBuilder):
         image_uri = config.get("destination")
         if image_uri is not None:
             registry = registry_from_uri(image_uri)
+
         return cls(
             environment,
             registry,
@@ -186,9 +187,8 @@ class KanikoBuilder(AbstractBuilder):
         Raises:
             LaunchError: If the builder config is invalid.
         """
-        if self.environment is None:
-            raise LaunchError("No environment specified for Kaniko build.")
-        await self.environment.verify_storage_uri(self.build_context_store)
+        if self.build_context_store:
+            await self.environment.verify_storage_uri(self.build_context_store)
 
     def login(self) -> None:
         """Login to the registry."""
@@ -197,8 +197,6 @@ class KanikoBuilder(AbstractBuilder):
     async def _create_docker_ecr_config_map(
         self, job_name: str, corev1_client: client.CoreV1Api, repository: str
     ) -> None:
-        if self.registry is None:
-            raise LaunchError("No registry specified for Kaniko build.")
         username, password = await self.registry.get_username_password()
         encoded = base64.b64encode(f"{username}:{password}".encode()).decode("utf-8")
         ecr_config_map = client.V1ConfigMap(
@@ -235,11 +233,21 @@ class KanikoBuilder(AbstractBuilder):
         with tarfile.TarFile.open(fileobj=context_file, mode="w:gz") as context_tgz:
             context_tgz.add(context_path, arcname=".")
         context_file.close()
-        destination = f"{self.build_context_store}/{run_id}.tgz"
-        if self.environment is None:
-            raise LaunchError("No environment specified for Kaniko build.")
-        await self.environment.upload_file(context_file.name, destination)
-        return destination
+        if PVC_MOUNT_PATH is None:
+            destination = f"{self.build_context_store}/{run_id}.tgz"
+            if self.environment is None:
+                raise LaunchError("No environment specified for Kaniko build.")
+            await self.environment.upload_file(context_file.name, destination)
+            return destination
+        else:
+            destination = f"{PVC_MOUNT_PATH}/{run_id}.tgz"
+            try:
+                shutil.copy(context_file.name, destination)
+            except Exception as e:
+                raise LaunchError(
+                    f"Error copying build context to PVC mounted at {PVC_MOUNT_PATH}: {e}"
+                ) from e
+            return f"tar:///context/{run_id}.tgz"
 
     async def build_image(
         self,
@@ -248,9 +256,6 @@ class KanikoBuilder(AbstractBuilder):
         job_tracker: Optional[JobAndRunStatusTracker] = None,
     ) -> str:
         await self.verify()
-        # TODO: this should probably throw an error if the registry is a local registry
-        if not self.registry:
-            raise LaunchError("No registry specified for Kaniko build.")
         # kaniko builder doesn't seem to work with a custom user id, need more investigation
         dockerfile_str = generate_dockerfile(
             launch_project=launch_project,
@@ -262,7 +267,6 @@ class KanikoBuilder(AbstractBuilder):
         image_tag = image_tag_from_dockerfile_and_source(launch_project, dockerfile_str)
         repo_uri = await self.registry.get_repo_uri()
         image_uri = repo_uri + ":" + image_tag
-
         if (
             not launch_project.build_required()
             and await self.registry.check_image_exists(image_uri)
@@ -271,16 +275,6 @@ class KanikoBuilder(AbstractBuilder):
 
         _logger.info(f"Building image {image_uri}...")
 
-        entry_cmd = " ".join(
-            get_entry_point_command(entrypoint, launch_project.override_args)
-        )
-
-        create_metadata_file(
-            launch_project,
-            image_uri,
-            sanitize_wandb_api_key(entry_cmd),
-            sanitize_wandb_api_key(dockerfile_str),
-        )
         context_path = _create_docker_build_ctx(launch_project, dockerfile_str)
         run_id = launch_project.run_id
 
@@ -381,6 +375,20 @@ class KanikoBuilder(AbstractBuilder):
         env = []
         volume_mounts = []
         volumes = []
+
+        if PVC_MOUNT_PATH:
+            volumes.append(
+                client.V1Volume(
+                    name="kaniko-pvc",
+                    persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource(
+                        claim_name=PVC_NAME
+                    ),
+                )
+            )
+            volume_mounts.append(
+                client.V1VolumeMount(name="kaniko-pvc", mount_path="/context")
+            )
+
         if bool(self.secret_name) != bool(self.secret_key):
             raise LaunchError(
                 "Both secret_name and secret_key or neither must be specified "
@@ -419,8 +427,27 @@ class KanikoBuilder(AbstractBuilder):
                     ),
                 )
             ]
-
-        if self.secret_name and self.secret_key:
+        if DOCKER_CONFIG_SECRET:
+            volumes.append(
+                client.V1Volume(
+                    name="kaniko-docker-config",
+                    secret=client.V1SecretVolumeSource(
+                        secret_name=DOCKER_CONFIG_SECRET,
+                        items=[
+                            client.V1KeyToPath(
+                                key=".dockerconfigjson", path="config.json"
+                            )
+                        ],
+                    ),
+                )
+            )
+            volume_mounts.append(
+                client.V1VolumeMount(
+                    name="kaniko-docker-config",
+                    mount_path="/kaniko/.docker",
+                )
+            )
+        elif self.secret_name and self.secret_key:
             volumes += [
                 client.V1Volume(
                     name="docker-config",

wandb/sdk/launch/create_job.py

@@ -315,39 +315,6 @@ def _create_repo_metadata(
         wandb.termerror(f"Entrypoint {entrypoint} not found in git repo")
         return None
 
-    # check if requirements.txt exists
-    # start at the location of the python file and recurse up to the git root
-    entrypoint_dir = os.path.dirname(entrypoint)
-    if entrypoint_dir:
-        req_dir = os.path.join(local_dir, entrypoint_dir)
-    else:
-        req_dir = local_dir
-
-    # If there is a Dockerfile.wandb in the starting rec dir, don't require a requirements.txt
-    if os.path.exists(os.path.join(req_dir, "Dockerfile.wandb")):
-        wandb.termlog(
-            f"Using Dockerfile.wandb in {req_dir.replace(tempdir, '') or 'repository root'}"
-        )
-    else:
-        while (
-            not os.path.exists(os.path.join(req_dir, "requirements.txt"))
-            and req_dir != tempdir
-        ):
-            req_dir = os.path.dirname(req_dir)
-
-        if not os.path.exists(os.path.join(req_dir, "requirements.txt")):
-            path_with_subdir = os.path.dirname(
-                os.path.join(path or "", entrypoint or "")
-            )
-            wandb.termerror(
-                f"Could not find requirements.txt file in git repo at {path_with_subdir}"
-            )
-            return None
-
-        wandb.termlog(
-            f"Using requirements.txt in {req_dir.replace(tempdir, '') or 'repository root'}"
-        )
-
     metadata = {
         "git": {
             "commit": commit,
@@ -366,18 +333,16 @@ def _create_repo_metadata(
 def _create_artifact_metadata(
     path: str, entrypoint: str, runtime: Optional[str] = None
 ) -> Tuple[Dict[str, Any], List[str]]:
-    if not os.path.exists(path):
+    if not os.path.isdir(path):
         wandb.termerror("Path must be a valid file or directory")
         return {}, []
 
-    if not os.path.exists(os.path.join(path, "requirements.txt")):
-        wandb.termerror(f"Could not find requirements.txt file in: {path}")
-        return {}, []
-
     # read local requirements.txt and dump to temp dir for builder
     requirements = []
-    with open(os.path.join(path, "requirements.txt")) as f:
-        requirements = f.read().splitlines()
+    depspath = os.path.join(path, "requirements.txt")
+    if os.path.exists(depspath):
+        with open(depspath) as f:
+            requirements = f.read().splitlines()
 
     if runtime:
         python_version = _clean_python_version(runtime)
@@ -431,6 +396,7 @@ def _configure_job_builder_for_partial(tmpdir: str, job_source: str) -> JobBuild
     settings.update({"files_dir": tmpdir, "job_source": job_source})
     job_builder = JobBuilder(
         settings=settings,  # type: ignore
+        verbose=True,
     )
     # never allow notebook runs
     job_builder._is_notebook_run = False

wandb/sdk/launch/loader.py

@@ -1,6 +1,7 @@
 """Utilities for the agent."""
 from typing import Any, Dict, Optional
 
+import wandb
 from wandb.apis.internal import Api
 from wandb.docker import is_docker_installed
 from wandb.sdk.launch.errors import LaunchError
@@ -87,6 +88,15 @@ def registry_from_config(
         from .registry.local_registry import LocalRegistry
 
         return LocalRegistry()  # This is the default, dummy registry.
+
+    wandb.termwarn(
+        "The `registry` block of the launch agent config is being deprecated. "
+        "Please specify an image repository URI under the `builder.destination` "
+        "key of your launch agent config. See "
+        "https://docs.wandb.ai/guides/launch/setup-agent-advanced#agent-configuration "
+        "for more information."
+    )
+
     registry_type = config.get("type")
     if registry_type is None or registry_type == "local":
         from .registry.local_registry import LocalRegistry

wandb/sdk/launch/registry/anon.py

@@ -0,0 +1,29 @@
+from typing import Tuple
+
+from wandb.docker import is_docker_installed
+from wandb.sdk.launch.utils import docker_image_exists
+
+from .abstract import AbstractRegistry
+
+
+class AnonynmousRegistry(AbstractRegistry):
+    def __init__(self, uri: str) -> None:
+        """Initialize the registry."""
+        self.uri = uri
+
+    async def get_username_password(self) -> Tuple[str, str]:
+        """Get the username and password for the registry."""
+        raise NotImplementedError("Anonymous registry does not require authentication")
+
+    async def get_repo_uri(self) -> str:
+        return self.uri
+
+    async def check_image_exists(self, image_uri: str) -> bool:
+        """Check if an image exists in the registry."""
+        if not is_docker_installed():
+            return False
+        return docker_image_exists(image_uri)
+
+    @classmethod
+    def from_config(cls, config: dict) -> "AbstractRegistry":
+        return cls(uri=config["uri"])

wandb/sdk/launch/registry/local_registry.py

@@ -2,6 +2,7 @@
 import logging
 from typing import Tuple
 
+from wandb.docker import is_docker_installed
 from wandb.sdk.launch.errors import LaunchError
 from wandb.sdk.launch.utils import docker_image_exists
 
@@ -60,4 +61,6 @@ class LocalRegistry(AbstractRegistry):
         Returns:
             bool: True.
         """
-        return docker_image_exists(image_uri)
+        if is_docker_installed():
+            return docker_image_exists(image_uri)
+        return False

wandb/sdk/launch/runner/kubernetes_runner.py

@@ -1,6 +1,7 @@
 """Implementation of KubernetesRunner class for wandb launch."""
 import asyncio
 import base64
+import datetime
 import json
 import logging
 import os
@@ -23,6 +24,7 @@ from wandb.sdk.launch.runner.kubernetes_monitor import (
     CustomResource,
     LaunchKubernetesMonitor,
 )
+from wandb.sdk.lib.retry import ExponentialBackoff, retry_async
 from wandb.util import get_module
 
 from .._project_spec import EntryPoint, LaunchProject
@@ -59,6 +61,7 @@ from kubernetes_asyncio.client.models.v1_secret import (  # type: ignore # noqa:
 from kubernetes_asyncio.client.rest import ApiException  # type: ignore # noqa: E402
 
 TIMEOUT = 5
+API_KEY_SECRET_MAX_RETRIES = 5
 
 _logger = logging.getLogger(__name__)
 
@@ -421,8 +424,23 @@ class KubernetesRunner(AbstractRunner):
                     else:
                         secret_name += f"-{launch_project.run_id}"
 
-                    api_key_secret = await ensure_api_key_secret(
-                        core_api, secret_name, namespace, value
+                    def handle_exception(e):
+                        wandb.termwarn(
+                            f"Exception when ensuring Kubernetes API key secret: {e}. Retrying..."
+                        )
+
+                    api_key_secret = await retry_async(
+                        backoff=ExponentialBackoff(
+                            initial_sleep=datetime.timedelta(seconds=1),
+                            max_sleep=datetime.timedelta(minutes=1),
+                            max_retries=API_KEY_SECRET_MAX_RETRIES,
+                        ),
+                        fn=ensure_api_key_secret,
+                        on_exc=handle_exception,
+                        core_api=core_api,
+                        secret_name=secret_name,
+                        namespace=namespace,
+                        api_key=value,
                     )
                     env.append(
                         {

wandb/sdk/launch/runner/local_container.py

@@ -148,15 +148,14 @@ class LocalContainerRunner(AbstractRunner):
             env_vars["WANDB_BASE_URL"] = "http://host.docker.internal:9001"
 
         if launch_project.docker_image:
-            if image_uri.endswith(":latest") or not docker_image_exists(image_uri):
-                try:
-                    pull_docker_image(image_uri)
-                except Exception as e:
-                    wandb.termwarn(f"Error attempting to pull docker image {image_uri}")
-                    if not docker_image_exists(image_uri):
-                        raise LaunchError(
-                            f"Failed to pull docker image {image_uri} with error: {e}"
-                        )
+            try:
+                pull_docker_image(image_uri)
+            except Exception as e:
+                wandb.termwarn(f"Error attempting to pull docker image {image_uri}")
+                if not docker_image_exists(image_uri):
+                    raise LaunchError(
+                        f"Failed to pull docker image {image_uri} with error: {e}"
+                    )
 
             assert launch_project.docker_image == image_uri
 
@@ -234,7 +233,13 @@ def _thread_process_runner(
         if not chunk:
             break
         index = chunk.find(b"\r")
-        decoded_chunk = chunk.decode()
+        decoded_chunk = None
+        while not decoded_chunk:
+            try:
+                decoded_chunk = chunk.decode()
+            except UnicodeDecodeError:
+                # Multi-byte character cut off, try to get the rest of it
+                chunk += os.read(process.stdout.fileno(), 1)  # type: ignore
         if index != -1:
             run._stdout += decoded_chunk
             print(chunk.decode(), end="")

wandb/sdk/launch/runner/sagemaker_runner.py

@@ -239,7 +239,7 @@ class SageMakerRunner(AbstractRunner):
             launch_project,
             self._api,
             role_arn,
-            launch_project.override_entrypoint,
+            entry_point,
             launch_project.override_args,
             MAX_ENV_LENGTHS[self.__class__.__name__],
             image_uri,

wandb/sdk/launch/sweeps/scheduler.py

@@ -1,6 +1,7 @@
 """Abstract Scheduler class."""
 import asyncio
 import base64
+import copy
 import logging
 import os
 import socket
@@ -24,7 +25,10 @@ from wandb.sdk.launch.sweeps.utils import (
     create_sweep_command_args,
     make_launch_sweep_entrypoint,
 )
-from wandb.sdk.launch.utils import event_loop_thread_exec
+from wandb.sdk.launch.utils import (
+    event_loop_thread_exec,
+    strip_resource_args_and_template_vars,
+)
 from wandb.sdk.lib.runid import generate_id
 
 if TYPE_CHECKING:
@@ -658,7 +662,7 @@ class Scheduler(ABC):
             pidx = entry_point.index("${program}")
             entry_point[pidx] = self._sweep_config["program"]
 
-        launch_config = self._wandb_run.config.get("launch", {})
+        launch_config = copy.deepcopy(self._wandb_run.config.get("launch", {}))
         if "overrides" not in launch_config:
             launch_config["overrides"] = {"run_config": {}}
         launch_config["overrides"]["run_config"].update(args["args_dict"])
@@ -694,11 +698,14 @@ class Scheduler(ABC):
             )
 
         # override resource and args of job
-        _job_launch_config = self._wandb_run.config.get("launch") or {}
+        _job_launch_config = copy.deepcopy(self._wandb_run.config.get("launch")) or {}
 
         # default priority is "medium"
         _priority = int(launch_config.get("priority", 2))  # type: ignore
 
+        # strip resource_args and template_variables from launch_config
+        strip_resource_args_and_template_vars(_job_launch_config)
+
         run_id = run.id or generate_id()
         queued_run = launch_add(
             run_id=run_id,
@@ -712,6 +719,7 @@ class Scheduler(ABC):
             project_queue=self._project_queue,
             resource=_job_launch_config.get("resource"),
             resource_args=_job_launch_config.get("resource_args"),
+            template_variables=_job_launch_config.get("template_variables"),
             author=self._kwargs.get("author"),
             sweep_id=self._sweep_id,
             priority=_priority,

wandb/sdk/launch/utils.py

@@ -221,6 +221,17 @@ def get_default_entity(api: Api, launch_config: Optional[Dict[str, Any]]):
     return config_entity or api.default_entity
 
 
+def strip_resource_args_and_template_vars(launch_spec: Dict[str, Any]) -> None:
+    if launch_spec.get("resource_args", None) and launch_spec.get(
+        "template_variables", None
+    ):
+        wandb.termwarn(
+            "Launch spec contains both resource_args and template_variables, "
+            "only one can be set. Using template_variables."
+        )
+        launch_spec.pop("resource_args")
+
+
 def construct_launch_spec(
     uri: Optional[str],
     job: Optional[str],
@@ -298,6 +309,9 @@ def construct_launch_spec(
         else:
             launch_config["registry"] = {"url": repository}
 
+    # dont send both resource args and template variables
+    strip_resource_args_and_template_vars(launch_spec)
+
     return launch_spec
 
 

wandb/sdk/lib/__init__.py

@@ -1,8 +1,5 @@
 from . import lazyloader
 from .disabled import RunDisabled, SummaryDisabled
+from .run_moment import RunMoment
 
-__all__ = (
-    "lazyloader",
-    "RunDisabled",
-    "SummaryDisabled",
-)
+__all__ = ("lazyloader", "RunDisabled", "SummaryDisabled", "RunMoment")

wandb/sdk/lib/_settings_toposort_generated.py

@@ -15,6 +15,7 @@ _Setting = Literal[
     "_aws_lambda",
     "_async_upload_concurrency_limit",
     "_cli_only_mode",
+    "_code_path_local",
     "_colab",
     "_cuda",
     "_disable_meta",
@@ -101,6 +102,7 @@ _Setting = Literal[
     "entity",
     "files_dir",
     "force",
+    "fork_from",
     "git_commit",
     "git_remote",
     "git_remote_url",
@@ -191,6 +193,8 @@ SETTINGS_TOPOLOGICALLY_SORTED: Final[Tuple[_Setting, ...]] = (
     "run_id",
     "start_method",
     "_aws_lambda",
+    "program",
+    "_code_path_local",
     "_colab",
     "_disable_machine_info",
     "_disable_meta",
@@ -227,7 +231,6 @@ SETTINGS_TOPOLOGICALLY_SORTED: Final[Tuple[_Setting, ...]] = (
     "log_symlink_internal",
     "log_symlink_user",
     "log_user",
-    "program",
     "project_url",
     "resume_fname",
     "run_url",

wandb/sdk/lib/apikey.py

@@ -179,11 +179,6 @@ def write_netrc(host: str, entity: str, key: str) -> Optional[bool]:
         return None
     try:
         normalized_host = urlparse(host).netloc.split(":")[0]
-        if normalized_host != "localhost" and "." not in normalized_host:
-            wandb.termerror(
-                f"Host must be a url in the form https://some.address.com, received {host}"
-            )
-            return None
         netrc_path = get_netrc_file_path()
         wandb.termlog(
             f"Appending key for {normalized_host} to your netrc file: {netrc_path}"