wandb 0.16.5__py3-none-any.whl → 0.17.0rc1__py3-none-any.whl

wandb/sdk/launch/builder/kaniko_builder.py

@@ -1,5 +1,6 @@
 import asyncio
 import base64
+import copy
 import json
 import logging
 import os
@@ -8,7 +9,7 @@ import tarfile
 import tempfile
 import time
 import traceback
-from typing import Optional
+from typing import Any, Dict, Optional
 
 import wandb
 from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
@@ -105,6 +106,7 @@ class KanikoBuilder(AbstractBuilder):
         secret_name: str = "",
         secret_key: str = "",
         image: str = "gcr.io/kaniko-project/executor:v1.11.0",
+        config: Optional[dict] = None,
     ):
         """Initialize a KanikoBuilder.
 
@@ -125,6 +127,7 @@ class KanikoBuilder(AbstractBuilder):
         self.secret_name = secret_name
         self.secret_key = secret_key
         self.image = image
+        self.kaniko_config = config or {}
 
     @classmethod
     def from_config(
@@ -170,6 +173,7 @@ class KanikoBuilder(AbstractBuilder):
         image_uri = config.get("destination")
         if image_uri is not None:
             registry = registry_from_uri(image_uri)
+        kaniko_config = config.get("kaniko-config", {})
 
         return cls(
             environment,
@@ -179,6 +183,7 @@ class KanikoBuilder(AbstractBuilder):
             secret_name=secret_name,
             secret_key=secret_key,
             image=kaniko_image,
+            config=kaniko_config,
         )
 
     async def verify(self) -> None:
@@ -289,7 +294,7 @@ class KanikoBuilder(AbstractBuilder):
 
         build_context = await self._upload_build_context(run_id, context_path)
         build_job = await self._create_kaniko_job(
-            build_job_name, repo_uri, image_uri, build_context, core_v1
+            build_job_name, repo_uri, image_uri, build_context, core_v1, api_client
         )
         wandb.termlog(f"{LOG_PREFIX}Created kaniko job {build_job_name}")
 
@@ -324,7 +329,9 @@ class KanikoBuilder(AbstractBuilder):
             ):
                 if job_tracker:
                     job_tracker.set_err_stage("build")
-                raise Exception(f"Failed to build image in kaniko for job {run_id}")
+                raise Exception(
+                    f"Failed to build image in kaniko for job {run_id}. View logs with `kubectl logs -n {NAMESPACE} {build_job_name}`."
+                )
             try:
                 pods_from_job = await core_v1.list_namespaced_pod(
                     namespace=NAMESPACE, label_selector=f"job-name={build_job_name}"
@@ -371,23 +378,32 @@ class KanikoBuilder(AbstractBuilder):
         image_tag: str,
         build_context_path: str,
         core_client: client.CoreV1Api,
-    ) -> "client.V1Job":
-        env = []
-        volume_mounts = []
-        volumes = []
+        api_client,
+    ) -> Dict[str, Any]:
+        job = copy.deepcopy(self.kaniko_config)
+        job_metadata = job.get("metadata", {})
+        job_labels = job_metadata.get("labels", {})
+        job_spec = job.get("spec", {})
+        pod_template = job_spec.get("template", {})
+        pod_metadata = pod_template.get("metadata", {})
+        pod_labels = pod_metadata.get("labels", {})
+        pod_spec = pod_template.get("spec", {})
+        volumes = pod_spec.get("volumes", [])
+        containers = pod_spec.get("containers") or [{}]
+        if len(containers) > 1:
+            raise LaunchError(
+                "Multiple container configs not supported for kaniko builder."
+            )
+        container = containers[0]
+        volume_mounts = container.get("volumeMounts", [])
+        env = container.get("env", [])
+        custom_args = container.get("args", [])
 
         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")
+                {"name": "kaniko-pvc", "persistentVolumeClaim": {"claimName": PVC_NAME}}
             )
+            volume_mounts.append({"name": "kaniko-pvc", "mountPath": "/context"})
 
         if bool(self.secret_name) != bool(self.secret_key):
             raise LaunchError(
@@ -395,13 +411,13 @@ class KanikoBuilder(AbstractBuilder):
                 "for kaniko build. You provided only one of them."
             )
         if isinstance(self.registry, ElasticContainerRegistry):
-            env += [
-                client.V1EnvVar(
-                    name="AWS_REGION",
-                    value=self.registry.region,
-                )
-            ]
-        # TODO: Refactor all of this environment/registry
+            env.append(
+                {
+                    "name": "AWS_REGION",
+                    "value": self.registry.region,
+                }
+            )
+        # TODO(ben): 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
@@ -416,52 +432,46 @@ class KanikoBuilder(AbstractBuilder):
                     "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",
-                        )
-                    ),
-                )
-            ]
+            env.append(
+                {
+                    "name": "AZURE_STORAGE_ACCESS_KEY",
+                    "valueFrom": {
+                        "secretKeyRef": {
+                            "name": "azure-storage-access-key",
+                            "key": "password",
+                        }
+                    },
+                }
+            )
         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"
-                            )
+                {
+                    "name": "kaniko-docker-config",
+                    "secret": {
+                        "secretName": DOCKER_CONFIG_SECRET,
+                        "items": [
+                            {
+                                "key": ".dockerconfigjson",
+                                "path": "config.json",
+                            }
                         ],
-                    ),
-                )
+                    },
+                }
             )
             volume_mounts.append(
-                client.V1VolumeMount(
-                    name="kaniko-docker-config",
-                    mount_path="/kaniko/.docker",
-                )
+                {"name": "kaniko-docker-config", "mountPath": "/kaniko/.docker"}
             )
         elif self.secret_name and self.secret_key:
-            volumes += [
-                client.V1Volume(
-                    name="docker-config",
-                    config_map=client.V1ConfigMapVolumeSource(
-                        name=f"docker-config-{job_name}",
-                    ),
-                ),
-            ]
-            volume_mounts += [
-                client.V1VolumeMount(
-                    name="docker-config", mount_path="/kaniko/.docker/"
-                ),
-            ]
-            # TODO: I don't like conditioning on the registry type here. As a
+            volumes.append(
+                {
+                    "name": "docker-config",
+                    "configMap": {"name": f"docker-config-{job_name}"},
+                }
+            )
+            volume_mounts.append(
+                {"name": "docker-config", "mountPath": "/kaniko/.docker"}
+            )
+            # TODO(ben): I don't like conditioning on the registry type here. As a
             # future change I want the registry and environment classes to provide
             # a list of environment variables and volume mounts that need to be
             # added to the job. The environment class provides credentials for
@@ -475,90 +485,95 @@ class KanikoBuilder(AbstractBuilder):
             elif isinstance(self.registry, GoogleArtifactRegistry):
                 mount_path = "/kaniko/.config/gcloud"
                 key = "config.json"
-                env += [
-                    client.V1EnvVar(
-                        name="GOOGLE_APPLICATION_CREDENTIALS",
-                        value="/kaniko/.config/gcloud/config.json",
-                    )
-                ]
+                env.append(
+                    {
+                        "name": "GOOGLE_APPLICATION_CREDENTIALS",
+                        "value": "/kaniko/.config/gcloud/config.json",
+                    }
+                )
             else:
                 raise LaunchError(
                     f"Registry type {type(self.registry)} not supported by kaniko"
                 )
-            volume_mounts += [
-                client.V1VolumeMount(
-                    name=self.secret_name,
-                    mount_path=mount_path,
-                    read_only=True,
-                )
-            ]
-            volumes += [
-                client.V1Volume(
-                    name=self.secret_name,
-                    secret=client.V1SecretVolumeSource(
-                        secret_name=self.secret_name,
-                        items=[client.V1KeyToPath(key=self.secret_key, path=key)],
-                    ),
-                )
-            ]
+            volumes.append(
+                {
+                    "name": self.secret_name,
+                    "secret": {
+                        "secretName": self.secret_name,
+                        "items": [{"key": self.secret_key, "path": key}],
+                    },
+                }
+            )
+            volume_mounts.append(
+                {
+                    "name": self.secret_name,
+                    "mountPath": mount_path,
+                    "readOnly": True,
+                }
+            )
         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}",
-                    ),
-                ),
-            ]
+            # Add the docker config map
+            volumes.append(
+                {
+                    "name": "docker-config",
+                    "configMap": {"name": f"docker-config-{job_name}"},
+                }
+            )
+            volume_mounts.append(
+                {"name": "docker-config", "mountPath": "/kaniko/.docker/"}
+            )
         # 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}",
-            f"--dockerfile={_WANDB_DOCKERFILE_NAME}",
-            f"--destination={destination}",
-            "--cache=true",
-            f"--cache-repo={repository.replace('https://', '')}",
-            "--snapshotMode=redo",
-            "--compressed-caching=false",
+        args = {
+            "--context": build_context_path,
+            "--dockerfile": _WANDB_DOCKERFILE_NAME,
+            "--destination": destination,
+            "--cache": "true",
+            "--cache-repo": repository.replace("https://", ""),
+            "--snapshot-mode": "redo",
+            "--compressed-caching": "false",
+        }
+        for custom_arg in custom_args:
+            arg_name, arg_value = custom_arg.split("=", 1)
+            args[arg_name] = arg_value
+        parsed_args = [
+            f"{arg_name}={arg_value}" for arg_name, arg_value in args.items()
         ]
-        container = client.V1Container(
-            name="wandb-container-build",
-            image=self.image,
-            args=args,
-            volume_mounts=volume_mounts,
-            env=env if env else None,
-        )
-        # Create and configure a spec section
-        labels = {"wandb": "launch"}
+        container["args"] = parsed_args
+
+        # Apply the rest of our defaults
+        pod_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=labels),
-            spec=client.V1PodSpec(
-                restart_policy="Never",
-                active_deadline_seconds=_DEFAULT_BUILD_TIMEOUT_SECS,
-                containers=[container],
-                volumes=volumes,
-                service_account_name=SERVICE_ACCOUNT_NAME,
-            ),
+            pod_labels["azure.workload.identity/use"] = "true"
+        pod_spec["restartPolicy"] = pod_spec.get("restartPolicy", "Never")
+        pod_spec["activeDeadlineSeconds"] = pod_spec.get(
+            "activeDeadlineSeconds", _DEFAULT_BUILD_TIMEOUT_SECS
         )
-        # Create the specification of job
-        spec = client.V1JobSpec(template=template, backoff_limit=0)
-        job = client.V1Job(
-            api_version="batch/v1",
-            kind="Job",
-            metadata=client.V1ObjectMeta(
-                name=job_name, namespace=NAMESPACE, labels={"wandb": "launch"}
-            ),
-            spec=spec,
+        pod_spec["serviceAccountName"] = pod_spec.get(
+            "serviceAccountName", SERVICE_ACCOUNT_NAME
         )
+        job_spec["backoffLimit"] = job_spec.get("backoffLimit", 0)
+        job_labels["wandb"] = "launch"
+        job_metadata["namespace"] = job_metadata.get("namespace", NAMESPACE)
+        job_metadata["name"] = job_metadata.get("name", job_name)
+        job["apiVersion"] = "batch/v1"
+        job["kind"] = "Job"
+
+        # Apply all nested configs from the bottom up
+        pod_metadata["labels"] = pod_labels
+        pod_template["metadata"] = pod_metadata
+        container["name"] = container.get("name", "wandb-container-build")
+        container["image"] = container.get("image", self.image)
+        container["volumeMounts"] = volume_mounts
+        container["env"] = env
+        pod_spec["containers"] = [container]
+        pod_spec["volumes"] = volumes
+        pod_template["spec"] = pod_spec
+        job_spec["template"] = pod_template
+        job_metadata["labels"] = job_labels
+        job["metadata"] = job_metadata
+        job["spec"] = job_spec
         return job

wandb/sdk/launch/builder/noop.py

@@ -1,4 +1,5 @@
 """NoOp builder implementation."""
+
 from typing import Any, Dict, Optional
 
 from wandb.sdk.launch.builder.abstract import AbstractBuilder

wandb/sdk/launch/create_job.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+import re
 import sys
 import tempfile
 from typing import Any, Dict, List, Optional, Tuple
@@ -11,7 +12,7 @@ from wandb.sdk.artifacts.artifact import Artifact
 from wandb.sdk.internal.job_builder import JobBuilder
 from wandb.sdk.launch.builder.build import get_current_python_version
 from wandb.sdk.launch.git_reference import GitReference
-from wandb.sdk.launch.utils import _is_git_uri
+from wandb.sdk.launch.utils import _is_git_uri, get_entrypoint_file
 from wandb.sdk.lib import filesystem
 from wandb.util import make_artifact_name_safe
 
@@ -19,6 +20,9 @@ logging.basicConfig(stream=sys.stdout, level=logging.INFO)
 _logger = logging.getLogger("wandb")
 
 
+CODE_ARTIFACT_EXCLUDE_PATHS = ["wandb", ".git"]
+
+
 def create_job(
     path: str,
     job_type: str,
@@ -107,6 +111,13 @@ def _create_job(
         )
         return None, "", []
 
+    if runtime is not None:
+        if not re.match(r"^3\.\d+$", runtime):
+            wandb.termerror(
+                f"Runtime (-r, --runtime) must be a minor version of Python 3, "
+                f"e.g. 3.9 or 3.10, received {runtime}"
+            )
+            return None, "", []
     aliases = aliases or []
     tempdir = tempfile.TemporaryDirectory()
     try:
@@ -145,6 +156,7 @@ def _create_job(
 
     job_builder = _configure_job_builder_for_partial(tempdir.name, job_source=job_type)
     if job_type == "code":
+        assert entrypoint is not None
         job_name = _make_code_artifact(
             api=api,
             job_builder=job_builder,
@@ -233,7 +245,6 @@ def _make_metadata_for_partial_job(
         return metadata, None
 
     if job_type == "code":
-        path, entrypoint = _handle_artifact_entrypoint(path, entrypoint)
         if not entrypoint:
             wandb.termerror(
                 "Artifact jobs must have an entrypoint, either included in the path or specified with -E"
@@ -304,15 +315,22 @@ def _create_repo_metadata(
             with open(os.path.join(local_dir, ".python-version")) as f:
                 python_version = f.read().strip().splitlines()[0]
         else:
-            major, minor = get_current_python_version()
-            python_version = f"{major}.{minor}"
+            _, python_version = get_current_python_version()
 
     python_version = _clean_python_version(python_version)
 
     # check if entrypoint is valid
     assert entrypoint is not None
-    if not os.path.exists(os.path.join(local_dir, entrypoint)):
-        wandb.termerror(f"Entrypoint {entrypoint} not found in git repo")
+    entrypoint_list = entrypoint.split(" ")
+    entrypoint_file = get_entrypoint_file(entrypoint_list)
+    if not entrypoint_file:
+        wandb.termerror(
+            f"Entrypoint {entrypoint} is invalid. An entrypoint should include both an executable and a file, for example 'python train.py'"
+        )
+        return None
+
+    if not os.path.exists(os.path.join(local_dir, entrypoint_file)):
+        wandb.termerror(f"Entrypoint file {entrypoint_file} not found in git repo")
         return None
 
     metadata = {
@@ -320,9 +338,9 @@ def _create_repo_metadata(
             "commit": commit,
             "remote": ref.url,
         },
-        "codePathLocal": entrypoint,  # not in git context, optionally also set local
-        "codePath": entrypoint,
-        "entrypoint": [f"python{python_version}", entrypoint],
+        "codePathLocal": entrypoint_file,  # not in git context, optionally also set local
+        "codePath": entrypoint_file,
+        "entrypoint": entrypoint_list,
         "python": python_version,  # used to build container
         "notebook": False,  # partial jobs from notebooks not supported
     }
@@ -332,10 +350,17 @@ def _create_repo_metadata(
 
 def _create_artifact_metadata(
     path: str, entrypoint: str, runtime: Optional[str] = None
-) -> Tuple[Dict[str, Any], List[str]]:
+) -> Tuple[Optional[Dict[str, Any]], Optional[List[str]]]:
     if not os.path.isdir(path):
         wandb.termerror("Path must be a valid file or directory")
         return {}, []
+    entrypoint_list = entrypoint.split(" ")
+    entrypoint_file = get_entrypoint_file(entrypoint_list)
+    if not entrypoint_file:
+        wandb.termerror(
+            f"Entrypoint {entrypoint} is invalid. An entrypoint should include both an executable and a file, for example 'python train.py'"
+        )
+        return None, None
 
     # read local requirements.txt and dump to temp dir for builder
     requirements = []
@@ -347,41 +372,17 @@ def _create_artifact_metadata(
     if runtime:
         python_version = _clean_python_version(runtime)
     else:
-        python_version = ".".join(get_current_python_version())
+        python_version, _ = get_current_python_version()
+        python_version = _clean_python_version(python_version)
 
-    metadata = {"python": python_version, "codePath": entrypoint}
+    metadata = {
+        "python": python_version,
+        "codePath": entrypoint_file,
+        "entrypoint": entrypoint_list,
+    }
     return metadata, requirements
 
 
-def _handle_artifact_entrypoint(
-    path: str, entrypoint: Optional[str] = None
-) -> Tuple[str, Optional[str]]:
-    if os.path.isfile(path):
-        if entrypoint and path.endswith(entrypoint):
-            path = path.replace(entrypoint, "")
-            wandb.termwarn(
-                f"Both entrypoint provided and path contains file. Using provided entrypoint: {entrypoint}, path is now: {path}"
-            )
-        elif entrypoint:
-            wandb.termwarn(
-                f"Ignoring passed in entrypoint as it does not match file path found in 'path'. Path entrypoint: {path.split('/')[-1]}"
-            )
-        entrypoint = path.split("/")[-1]
-        path = "/".join(path.split("/")[:-1])
-    elif not entrypoint:
-        wandb.termerror("Entrypoint not valid")
-        return "", None
-    path = path or "."  # when path is just an entrypoint, use cdw
-
-    if not os.path.exists(os.path.join(path, entrypoint)):
-        wandb.termerror(
-            f"Could not find execution point: {os.path.join(path, entrypoint)}"
-        )
-        return "", None
-
-    return path, entrypoint
-
-
 def _configure_job_builder_for_partial(tmpdir: str, job_source: str) -> JobBuilder:
     """Configure job builder with temp dir and job source."""
     # adjust git source to repo
@@ -411,7 +412,7 @@ def _make_code_artifact(
     job_builder: JobBuilder,
     run: "wandb.sdk.wandb_run.Run",
     path: str,
-    entrypoint: Optional[str],
+    entrypoint: str,
     entity: Optional[str],
     project: Optional[str],
     name: Optional[str],
@@ -420,17 +421,22 @@ def _make_code_artifact(
 
     Returns the name of the eventual job.
     """
-    artifact_name = _make_code_artifact_name(os.path.join(path, entrypoint or ""), name)
+    assert entrypoint is not None
+    entrypoint_list = entrypoint.split(" ")
+    entrypoint_file = get_entrypoint_file(entrypoint_list)
+    if not entrypoint_file:
+        wandb.termerror(
+            f"Entrypoint {entrypoint} is invalid. An entrypoint should include both an executable and a file, for example 'python train.py'"
+        )
+        return None
+
+    artifact_name = _make_code_artifact_name(os.path.join(path, entrypoint_file), name)
     code_artifact = wandb.Artifact(
         name=artifact_name,
         type="code",
         description="Code artifact for job",
     )
 
-    # Update path and entrypoint vars to match metadata
-    # TODO(gst): consolidate into one place
-    path, entrypoint = _handle_artifact_entrypoint(path, entrypoint)
-
     try:
         code_artifact.add_dir(path)
     except Exception as e:
@@ -441,6 +447,13 @@ def _make_code_artifact(
         wandb.termerror(f"Error adding to code artifact: {e}")
         return None
 
+    # Remove paths we don't want to include, if present
+    for item in CODE_ARTIFACT_EXCLUDE_PATHS:
+        try:
+            code_artifact.remove(item)
+        except FileNotFoundError:
+            pass
+
     res, _ = api.create_artifact(
         artifact_type_name="code",
         artifact_collection_name=artifact_name,
@@ -451,7 +464,7 @@ def _make_code_artifact(
         project_name=project,
         run_name=run.id,  # run will be deleted after creation
         description="Code artifact for job",
-        metadata={"codePath": path, "entrypoint": entrypoint},
+        metadata={"codePath": path, "entrypoint": entrypoint_file},
         is_user_created=True,
         aliases=[
             {"artifactCollectionName": artifact_name, "alias": a} for a in ["latest"]

wandb/sdk/launch/environment/abstract.py

@@ -1,4 +1,5 @@
 """Abstract base class for environments."""
+
 from abc import ABC, abstractmethod
 
 

wandb/sdk/launch/environment/gcp_environment.py

@@ -1,4 +1,5 @@
 """Implementation of the GCP environment for wandb launch."""
+
 import logging
 import os
 import subprocess

wandb/sdk/launch/environment/local_environment.py

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

wandb/sdk/launch/loader.py

@@ -1,4 +1,5 @@
 """Utilities for the agent."""
+
 from typing import Any, Dict, Optional
 
 import wandb

wandb/sdk/launch/registry/abstract.py

@@ -1,4 +1,5 @@
 """Abstract base class for registries."""
+
 from abc import ABC, abstractmethod
 from typing import Tuple
 

wandb/sdk/launch/registry/azure_container_registry.py

@@ -1,4 +1,5 @@
 """Implementation of AzureContainerRegistry class."""
+
 import re
 from typing import TYPE_CHECKING, Optional, Tuple
 

wandb/sdk/launch/registry/elastic_container_registry.py

@@ -1,4 +1,5 @@
 """Implementation of Elastic Container Registry class for wandb launch."""
+
 import base64
 import logging
 from typing import Dict, Optional, Tuple

wandb/sdk/launch/registry/google_artifact_registry.py

@@ -1,4 +1,5 @@
 """Implementation of Google Artifact Registry for wandb launch."""
+
 import logging
 from typing import Optional, Tuple