wandb 0.17.5__py3-none-any.whl → 0.17.6__py3-none-any.whl

wandb/sdk/launch/_project_spec.py

@@ -7,6 +7,7 @@ import enum
 import json
 import logging
 import os
+import shutil
 import tempfile
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
@@ -112,6 +113,9 @@ class LaunchProject:
         self.sweep_id = sweep_id
         self.author = launch_spec.get("author")
         self.python_version: Optional[str] = launch_spec.get("python_version")
+        self._job_dockerfile: Optional[str] = None
+        self._job_build_context: Optional[str] = None
+        self._job_base_image: Optional[str] = None
         self.accelerator_base_image: Optional[str] = resource_args_build.get(
             "accelerator", {}
         ).get("base_image") or resource_args_build.get("cuda", {}).get("base_image")
@@ -131,8 +135,6 @@ class LaunchProject:
         self._queue_name: Optional[str] = None
         self._queue_entity: Optional[str] = None
         self._run_queue_item_id: Optional[str] = None
-        self._job_dockerfile: Optional[str] = None
-        self._job_build_context: Optional[str] = None
 
     def init_source(self) -> None:
         if self.docker_image is not None:
@@ -146,6 +148,21 @@ class LaunchProject:
             self.project_dir = os.getcwd()
             self._entry_point = self.override_entrypoint
 
+    def change_project_dir(self, new_dir: str) -> None:
+        """Change the project directory to a new directory."""
+        # Copy the contents of the old project dir to the new project dir.
+        old_dir = self.project_dir
+        if old_dir is not None:
+            shutil.copytree(
+                old_dir,
+                new_dir,
+                symlinks=True,
+                dirs_exist_ok=True,
+                ignore=shutil.ignore_patterns("fsmonitor--daemon.ipc", ".git"),
+            )
+            shutil.rmtree(old_dir)
+        self.project_dir = new_dir
+
     def init_git(self, git_info: Dict[str, str]) -> None:
         self.git_version = git_info.get("version")
         self.git_repo = git_info.get("repo")
@@ -212,14 +229,23 @@ class LaunchProject:
     def job_build_context(self) -> Optional[str]:
         return self._job_build_context
 
+    @property
+    def job_base_image(self) -> Optional[str]:
+        return self._job_base_image
+
     def set_job_dockerfile(self, dockerfile: str) -> None:
         self._job_dockerfile = dockerfile
 
     def set_job_build_context(self, build_context: str) -> None:
         self._job_build_context = build_context
 
+    def set_job_base_image(self, base_image: str) -> None:
+        self._job_base_image = base_image
+
     @property
     def image_name(self) -> str:
+        if self.job_base_image is not None:
+            return self.job_base_image
         if self.docker_image is not None:
             return self.docker_image
         elif self.uri is not None:
@@ -299,10 +325,8 @@ class LaunchProject:
 
     def build_required(self) -> bool:
         """Checks the source to see if a build is required."""
-        # since the image tag for images built from jobs
-        # is based on the job version index, which is immutable
-        # we don't need to build the image for a job if that tag
-        # already exists
+        if self.job_base_image is not None:
+            return False
         if self.source != LaunchSource.JOB:
             return True
         return False
@@ -316,7 +340,9 @@ class LaunchProject:
         Returns:
             Optional[str]: The Docker image or None if not specified.
         """
-        return self._docker_image
+        if self._docker_image:
+            return self._docker_image
+        return None
 
     @docker_image.setter
     def docker_image(self, value: str) -> None:
@@ -336,7 +362,7 @@ class LaunchProject:
         # assuming project only has 1 entry point, pull that out
         # tmp fn until we figure out if we want to support multiple entry points or not
         if not self._entry_point:
-            if not self.docker_image:
+            if not self.docker_image and not self.job_base_image:
                 raise LaunchError(
                     "Project must have at least one entry point unless docker image is specified."
                 )

wandb/sdk/launch/agent/agent.py

@@ -717,7 +717,7 @@ class LaunchAgent:
         _, build_config, registry_config = construct_agent_configs(
             default_config, override_build_config
         )
-        image_uri = project.docker_image
+        image_uri = project.docker_image or project.job_base_image
         entrypoint = project.get_job_entry_point()
         environment = loader.environment_from_config(
             default_config.get("environment", {})
@@ -727,7 +727,11 @@ class LaunchAgent:
         backend = loader.runner_from_config(
             resource, api, backend_config, environment, registry
         )
-        if not (project.docker_image or isinstance(backend, LocalProcessRunner)):
+        if not (
+            project.docker_image
+            or project.job_base_image
+            or isinstance(backend, LocalProcessRunner)
+        ):
             assert entrypoint is not None
             image_uri = await builder.build_image(project, entrypoint, job_tracker)
 

wandb/sdk/launch/agent/run_queue_item_file_saver.py

@@ -2,7 +2,7 @@
 
 import os
 import sys
-from typing import List, Optional, Union
+from typing import List, Optional
 
 import wandb
 
@@ -17,9 +17,7 @@ FileSubtypes = Literal["warning", "error"]
 class RunQueueItemFileSaver:
     def __init__(
         self,
-        agent_run: Optional[
-            Union["wandb.sdk.wandb_run.Run", "wandb.sdk.lib.RunDisabled"]
-        ],
+        agent_run: Optional["wandb.sdk.wandb_run.Run"],
         run_queue_item_id: str,
     ):
         self.run_queue_item_id = run_queue_item_id

wandb/sdk/launch/builder/build.py

@@ -201,7 +201,7 @@ def get_requirements_section(
     # If there is a requirements.txt at root of build context, use that.
     if (base_path / "src" / "requirements.txt").exists():
         requirements_files += ["src/requirements.txt"]
-        deps_install_line = "pip install -r requirements.txt"
+        deps_install_line = "pip install uv && uv pip install -r requirements.txt"
         with open(base_path / "src" / "requirements.txt") as f:
             requirements = f.readlines()
         if not any(["wandb" in r for r in requirements]):
@@ -237,7 +237,9 @@ def get_requirements_section(
                 with open(base_path / "src" / "requirements.txt", "w") as f:
                     f.write("\n".join(project_deps))
                 requirements_files += ["src/requirements.txt"]
-                deps_install_line = "pip install -r requirements.txt"
+                deps_install_line = (
+                    "pip install uv && uv pip install -r requirements.txt"
+                )
                 return PIP_TEMPLATE.format(
                     buildx_optional_prefix=prefix,
                     requirements_files=" ".join(requirements_files),

wandb/sdk/launch/builder/templates/_wandb_bootstrap.py

@@ -39,12 +39,13 @@ def install_deps(
         deps (str[], None): The dependencies that failed to install
     """
     try:
+        subprocess.check_output(["pip", "install", "uv"], stderr=subprocess.STDOUT)
         # Include only uri if @ is present
         clean_deps = [d.split("@")[-1].strip() if "@" in d else d for d in deps]
         index_args = ["--extra-index-url", extra_index] if extra_index else []
         print("installing {}...".format(", ".join(clean_deps)))
         opts = opts or []
-        args = ["pip", "install"] + opts + clean_deps + index_args
+        args = ["uv", "pip", "install"] + opts + clean_deps + index_args
         sys.stdout.flush()
         subprocess.check_output(args, stderr=subprocess.STDOUT)
         return failed

wandb/sdk/launch/inputs/internal.py

@@ -16,7 +16,9 @@ from typing import Any, Dict, List, Optional
 import wandb
 import wandb.data_types
 from wandb.sdk.launch.errors import LaunchError
+from wandb.sdk.launch.inputs.schema import META_SCHEMA
 from wandb.sdk.wandb_run import Run
+from wandb.util import get_module
 
 from .files import config_path_is_valid, override_file
 
@@ -129,7 +131,7 @@ def _publish_job_input(
     )
 
 
-def _replace_refs_and_allofs(schema: dict, defs: dict) -> dict:
+def _replace_refs_and_allofs(schema: dict, defs: Optional[dict]) -> dict:
     """Recursively fix JSON schemas with common issues.
 
     1. Replaces any instances of $ref with their associated definition in defs
@@ -137,7 +139,7 @@ def _replace_refs_and_allofs(schema: dict, defs: dict) -> dict:
     See test_internal.py for examples
     """
     ret: Dict[str, Any] = {}
-    if "$ref" in schema:
+    if "$ref" in schema and defs:
         # Reference found, replace it with its definition
         def_key = schema["$ref"].split("#/$defs/")[1]
         # Also run recursive replacement in case a ref contains more refs
@@ -170,12 +172,16 @@ def _replace_refs_and_allofs(schema: dict, defs: dict) -> dict:
     return ret
 
 
-def _convert_pydantic_model_to_jsonschema(model: Any) -> dict:
-    schema = model.model_json_schema()
-    defs = schema.pop("$defs")
-    if not defs:
-        return schema
-    return _replace_refs_and_allofs(schema, defs)
+def _validate_schema(schema: dict) -> None:
+    jsonschema = get_module(
+        "jsonschema",
+        required="Setting job schema requires the jsonschema package. Please install it with `pip install 'wandb[launch]'`.",
+        lazy=False,
+    )
+    validator = jsonschema.Draft202012Validator(META_SCHEMA)
+    errs = sorted(validator.iter_errors(schema), key=str)
+    if errs:
+        wandb.termwarn(f"Schema includes unhandled or invalid configurations:\n{errs}")
 
 
 def handle_config_file_input(
@@ -204,16 +210,20 @@ def handle_config_file_input(
         path,
         dest,
     )
-    # This supports both an instance of a pydantic BaseModel class (e.g. schema=MySchema(...))
-    # or the BaseModel class itself (e.g. schema=MySchema)
-    if hasattr(schema, "model_json_schema") and callable(
-        schema.model_json_schema  # type: ignore
-    ):
-        schema = _convert_pydantic_model_to_jsonschema(schema)
-    if schema and not isinstance(schema, dict):
-        raise LaunchError(
-            "schema must be a dict, Pydantic model instance, or Pydantic model class."
-        )
+    if schema:
+        # This supports both an instance of a pydantic BaseModel class (e.g. schema=MySchema(...))
+        # or the BaseModel class itself (e.g. schema=MySchema)
+        if hasattr(schema, "model_json_schema") and callable(
+            schema.model_json_schema  # type: ignore
+        ):
+            schema = schema.model_json_schema()
+        if not isinstance(schema, dict):
+            raise LaunchError(
+                "schema must be a dict, Pydantic model instance, or Pydantic model class."
+            )
+        defs = schema.pop("$defs", None)
+        schema = _replace_refs_and_allofs(schema, defs)
+        _validate_schema(schema)
     arguments = JobInputArguments(
         include=include,
         exclude=exclude,
@@ -241,16 +251,20 @@ def handle_run_config_input(
     If there is no active run, the include and exclude paths are staged and sent
     when a run is created.
     """
-    # This supports both an instance of a pydantic BaseModel class (e.g. schema=MySchema(...))
-    # or the BaseModel class itself (e.g. schema=MySchema)
-    if hasattr(schema, "model_json_schema") and callable(
-        schema.model_json_schema  # type: ignore
-    ):
-        schema = _convert_pydantic_model_to_jsonschema(schema)
-    if schema and not isinstance(schema, dict):
-        raise LaunchError(
-            "schema must be a dict, Pydantic model instance, or Pydantic model class."
-        )
+    if schema:
+        # This supports both an instance of a pydantic BaseModel class (e.g. schema=MySchema(...))
+        # or the BaseModel class itself (e.g. schema=MySchema)
+        if hasattr(schema, "model_json_schema") and callable(
+            schema.model_json_schema  # type: ignore
+        ):
+            schema = schema.model_json_schema()
+        if not isinstance(schema, dict):
+            raise LaunchError(
+                "schema must be a dict, Pydantic model instance, or Pydantic model class."
+            )
+        defs = schema.pop("$defs", None)
+        schema = _replace_refs_and_allofs(schema, defs)
+        _validate_schema(schema)
     arguments = JobInputArguments(
         include=include,
         exclude=exclude,

wandb/sdk/launch/inputs/schema.py

@@ -0,0 +1,39 @@
+META_SCHEMA = {
+    "type": "object",
+    "properties": {
+        "type": {
+            "type": "string",
+            "enum": ["boolean", "integer", "number", "string", "object"],
+        },
+        "title": {"type": "string"},
+        "description": {"type": "string"},
+        "enum": {"type": "array", "items": {"type": ["integer", "number", "string"]}},
+        "properties": {"type": "object", "patternProperties": {".*": {"$ref": "#"}}},
+        "allOf": {"type": "array", "items": {"$ref": "#"}},
+    },
+    "allOf": [
+        {
+            "if": {"properties": {"type": {"const": "number"}}},
+            "then": {
+                "properties": {
+                    "minimum": {"type": ["integer", "number"]},
+                    "maximum": {"type": ["integer", "number"]},
+                    "exclusiveMinimum": {"type": ["integer", "number"]},
+                    "exclusiveMaximum": {"type": ["integer", "number"]},
+                }
+            },
+        },
+        {
+            "if": {"properties": {"type": {"const": "integer"}}},
+            "then": {
+                "properties": {
+                    "minimum": {"type": "integer"},
+                    "maximum": {"type": "integer"},
+                    "exclusiveMinimum": {"type": "integer"},
+                    "exclusiveMaximum": {"type": "integer"},
+                }
+            },
+        },
+    ],
+    "unevaluatedProperties": False,
+}

wandb/sdk/launch/runner/kubernetes_runner.py

@@ -31,6 +31,7 @@ from wandb.util import get_module
 from .._project_spec import EntryPoint, LaunchProject
 from ..errors import LaunchError
 from ..utils import (
+    CODE_MOUNT_DIR,
     LOG_PREFIX,
     MAX_ENV_LENGTHS,
     PROJECT_SYNCHRONOUS,
@@ -66,6 +67,10 @@ API_KEY_SECRET_MAX_RETRIES = 5
 _logger = logging.getLogger(__name__)
 
 
+SOURCE_CODE_PVC_MOUNT_PATH = os.environ.get("WANDB_LAUNCH_CODE_PVC_MOUNT_PATH")
+SOURCE_CODE_PVC_NAME = os.environ.get("WANDB_LAUNCH_CODE_PVC_NAME")
+
+
 class KubernetesSubmittedRun(AbstractRun):
     """Wrapper for a launched run on Kubernetes."""
 
@@ -468,6 +473,12 @@ class KubernetesRunner(AbstractRunner):
             "true",
         )
 
+        if launch_project.job_base_image:
+            apply_code_mount_configuration(
+                job,
+                launch_project,
+            )
+
         # Add wandb.ai/agent: current agent label on all pods
         if LaunchAgent.initialized():
             add_label_to_pods(
@@ -504,6 +515,22 @@ class KubernetesRunner(AbstractRunner):
             kubernetes_asyncio, resource_args
         )
 
+        # If using pvc for code mount, move code there.
+        if launch_project.job_base_image is not None:
+            if SOURCE_CODE_PVC_NAME is None or SOURCE_CODE_PVC_MOUNT_PATH is None:
+                raise LaunchError(
+                    "WANDB_LAUNCH_SOURCE_CODE_PVC_ environment variables not set. "
+                    "Unable to mount source code PVC into base image. "
+                    "Use the `codeMountPvcName` variable in the agent helm chart "
+                    "to enable base image jobs for this agent. See "
+                    "https://github.com/wandb/helm-charts/tree/main/charts/launch-agent "
+                    "for more information."
+                )
+            code_subdir = launch_project.get_image_source_string()
+            launch_project.change_project_dir(
+                os.path.join(SOURCE_CODE_PVC_MOUNT_PATH, code_subdir)
+            )
+
         # If the user specified an alternate api, we need will execute this
         # run by creating a custom object.
         api_version = resource_args.get("apiVersion", "batch/v1")
@@ -542,6 +569,9 @@ class KubernetesRunner(AbstractRunner):
                     LaunchAgent.name()
                 )
 
+            if launch_project.job_base_image:
+                apply_code_mount_configuration(resource_args, launch_project)
+
             overrides = {}
             if launch_project.override_args:
                 overrides["args"] = launch_project.override_args
@@ -889,3 +919,45 @@ def add_entrypoint_args_overrides(manifest: Union[dict, list], overrides: dict)
                     container["args"] = overrides["args"]
         for value in manifest.values():
             add_entrypoint_args_overrides(value, overrides)
+
+
+def apply_code_mount_configuration(
+    manifest: Union[Dict, list], project: LaunchProject
+) -> None:
+    """Apply code mount configuration to all containers in a manifest.
+
+    Recursively traverses the manifest and adds the code mount configuration to
+    all containers. Containers are identified by the presence of a "spec" key
+    with a "containers" key in the value.
+
+    Arguments:
+        manifest: The manifest to modify.
+        project: The launch project.
+
+    Returns: None.
+    """
+    assert SOURCE_CODE_PVC_NAME is not None
+    source_dir = project.get_image_source_string()
+    for pod in yield_pods(manifest):
+        for container in yield_containers(pod):
+            if "volumeMounts" not in container:
+                container["volumeMounts"] = []
+            container["volumeMounts"].append(
+                {
+                    "name": "wandb-source-code-volume",
+                    "mountPath": CODE_MOUNT_DIR,
+                    "subPath": source_dir,
+                }
+            )
+            container["workingDir"] = CODE_MOUNT_DIR
+        spec = pod["spec"]
+        if "volumes" not in spec:
+            spec["volumes"] = []
+        spec["volumes"].append(
+            {
+                "name": "wandb-source-code-volume",
+                "persistentVolumeClaim": {
+                    "claimName": SOURCE_CODE_PVC_NAME,
+                },
+            }
+        )

wandb/sdk/launch/runner/local_container.py

@@ -14,6 +14,7 @@ from wandb.sdk.launch.registry.abstract import AbstractRegistry
 from .._project_spec import LaunchProject
 from ..errors import LaunchError
 from ..utils import (
+    CODE_MOUNT_DIR,
     LOG_PREFIX,
     MAX_ENV_LENGTHS,
     PROJECT_SYNCHRONOUS,
@@ -121,7 +122,15 @@ class LocalContainerRunner(AbstractRunner):
                 docker_args["network"] = "host"
             if sys.platform == "linux" or sys.platform == "linux2":
                 docker_args["add-host"] = "host.docker.internal:host-gateway"
-
+        base_image = launch_project.job_base_image
+        if base_image is not None:
+            # Mount code into the container and set the working directory.
+            if "volume" not in docker_args:
+                docker_args["volume"] = []
+            docker_args["volume"].append(
+                f"{launch_project.project_dir}:{CODE_MOUNT_DIR}"
+            )
+            docker_args["workdir"] = CODE_MOUNT_DIR
         return docker_args
 
     async def run(
@@ -146,7 +155,7 @@ class LocalContainerRunner(AbstractRunner):
         elif _is_wandb_dev_uri(self._api.settings("base_url")):
             env_vars["WANDB_BASE_URL"] = "http://host.docker.internal:9001"
 
-        if launch_project.docker_image:
+        if launch_project.docker_image or launch_project.job_base_image:
             try:
                 pull_docker_image(image_uri)
             except Exception as e:
@@ -156,14 +165,8 @@ class LocalContainerRunner(AbstractRunner):
                         f"Failed to pull docker image {image_uri} with error: {e}"
                     )
 
-            assert launch_project.docker_image == image_uri
-
-        entry_cmd = (
-            launch_project.override_entrypoint.command
-            if launch_project.override_entrypoint is not None
-            else None
-        )
-
+        entrypoint = launch_project.get_job_entry_point()
+        entry_cmd = None if entrypoint is None else entrypoint.command
         command_str = " ".join(
             get_docker_command(
                 image_uri,

wandb/sdk/launch/runner/sagemaker_runner.py

@@ -221,7 +221,6 @@ class SageMakerRunner(AbstractRunner):
                 await run.wait()
             return run
 
-        launch_project.fill_macros(image_uri)
         _logger.info("Connecting to sagemaker client")
         entry_point = (
             launch_project.override_entrypoint or launch_project.get_job_entry_point()
@@ -296,13 +295,12 @@ def build_sagemaker_args(
     entry_point: Optional[EntryPoint],
     args: Optional[List[str]],
     max_env_length: int,
-    image_uri: Optional[str] = None,
+    image_uri: str,
     default_output_path: Optional[str] = None,
 ) -> Dict[str, Any]:
     sagemaker_args: Dict[str, Any] = {}
-    given_sagemaker_args: Optional[Dict[str, Any]] = launch_project.resource_args.get(
-        "sagemaker"
-    )
+    resource_args = launch_project.fill_macros(image_uri)
+    given_sagemaker_args: Optional[Dict[str, Any]] = resource_args.get("sagemaker")
 
     if given_sagemaker_args is None:
         raise LaunchError(

wandb/sdk/launch/utils.py

@@ -87,6 +87,8 @@ LOG_PREFIX = f"{click.style('launch:', fg='magenta')} "
 MAX_ENV_LENGTHS: Dict[str, int] = defaultdict(lambda: 32670)
 MAX_ENV_LENGTHS["SageMakerRunner"] = 512
 
+CODE_MOUNT_DIR = "/mnt/wandb"
+
 
 def load_wandb_config() -> Config:
     """Load wandb config from WANDB_CONFIG environment variable(s).