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

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).

wandb/sdk/lib/apikey.py

@@ -107,7 +107,7 @@ def prompt_api_key(  # noqa: C901
 
     if jupyter and "google.colab" in sys.modules:
         log_string = term.LOG_STRING_NOCOLOR
-        key = wandb.jupyter.attempt_colab_login(app_url)
+        key = wandb.jupyter.attempt_colab_login(app_url)  # type: ignore
         if key is not None:
             write_key(settings, key, api=api)
             return key  # type: ignore

wandb/sdk/lib/disabled.py

@@ -1,178 +1,6 @@
-#
+from typing import Any
 
-
-class RunDisabled(str):
-    def __init__(self, *args, **kwargs):
-        object.__setattr__(self, "___dict", {})
-
-    def __add__(self, other):
-        return self
-
-    def __sub__(self, other):
-        return self
-
-    def __mul__(self, other):
-        return self
-
-    def __truediv__(self, other):
-        return self
-
-    def __floordiv__(self, other):
-        return self
-
-    def __mod__(self, other):
-        return self
-
-    def __pow__(self, other, modulo=None):
-        return self
-
-    def __lshift__(self, other):
-        return self
-
-    def __rshift__(self, other):
-        return self
-
-    def __and__(self, other):
-        return self
-
-    def __xor__(self, other):
-        return self
-
-    def __or__(self, other):
-        return self
-
-    def __iadd__(self, other):
-        return self
-
-    def __isub__(self, other):
-        return self
-
-    def __imul__(self, other):
-        return self
-
-    def __idiv__(self, other):
-        return self
-
-    def __ifloordiv__(self, other):
-        return self
-
-    def __imod__(self, other):
-        return self
-
-    def __ipow__(self, other, modulo=None):
-        return self
-
-    def __ilshift__(self, other):
-        return self
-
-    def __irshift__(self, other):
-        return self
-
-    def __iand__(self, other):
-        return self
-
-    def __ixor__(self, other):
-        return self
-
-    def __ior__(self, other):
-        return self
-
-    def __neg__(self):
-        return self
-
-    def __pos__(self):
-        return self
-
-    def __abs__(self):
-        return self
-
-    def __invert__(self):
-        return self
-
-    def __complex__(self):
-        return 1 + 0j
-
-    def __int__(self):
-        return 1
-
-    def __long__(self):
-        return 1
-
-    def __float__(self):
-        return 1.0
-
-    def __oct__(self):
-        return oct(1)
-
-    def __hex__(self):
-        return hex(1)
-
-    def __lt__(self, other):
-        return True
-
-    def __le__(self, other):
-        return True
-
-    def __eq__(self, other):
-        return True
-
-    def __ne__(self, other):
-        return True
-
-    def __gt__(self, other):
-        return True
-
-    def __ge__(self, other):
-        return True
-
-    def __getattr__(self, attr):
-        return self[attr]
-
-    def __getitem__(self, key):
-        d = object.__getattribute__(self, "___dict")
-        try:
-            if key in d:
-                return d[key]
-        except TypeError:
-            key = str(key)
-            if key in d:
-                return d[key]
-        dummy = RunDisabled()
-        d[key] = dummy
-        return dummy
-
-    def __setitem__(self, key, value):
-        object.__getattribute__(self, "___dict")[key] = value
-
-    def __setattr__(self, key, value):
-        self[key] = value
-
-    def __call__(self, *args, **kwargs):
-        return RunDisabled()
-
-    def __len__(self):
-        return 1
-
-    def __str__(self):
-        return ""
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        return exc_type is None
-
-    def __repr__(self):
-        return ""
-
-    def __nonzero__(self):
-        return True
-
-    def __bool__(self):
-        return True
-
-    def __getstate__(self):
-        return 1
+from wandb.sdk.lib import deprecate
 
 
 class SummaryDisabled(dict):
@@ -188,3 +16,14 @@ class SummaryDisabled(dict):
             val = SummaryDisabled(val)
             self[key] = val
         return val
+
+
+class RunDisabled:
+    """Compatibility class for integrations that explicitly check for wandb.RunDisabled."""
+
+    def __getattr__(self, name: str) -> Any:
+        deprecate.deprecate(
+            field_name=deprecate.Deprecated.run_disabled,
+            warning_message="RunDisabled is deprecated and is a no-op. "
+            '`wandb.init(mode="disabled")` now returns and instance of `wandb.sdk.wandb_run.Run`.',
+        )

wandb/sdk/service/streams.py

@@ -81,9 +81,7 @@ class StreamRecord:
         self._wait_thread_active()
 
     def _wait_thread_active(self) -> None:
-        result = self._iface.communicate_status()
-        # TODO: using the default communicate timeout, is that enough? retries?
-        assert result
+        self._iface.deliver_status().wait(timeout=-1)
 
     def join(self) -> None:
         self._iface.join()
@@ -212,7 +210,7 @@ class StreamMux:
         # run_id = action.stream_id  # will want to fix if a streamid != runid
         settings = action._data
         thread = StreamThread(
-            target=wandb.wandb_sdk.internal.internal.wandb_internal,
+            target=wandb.wandb_sdk.internal.internal.wandb_internal,  # type: ignore
             kwargs=dict(
                 settings=settings,
                 record_q=stream._record_q,

wandb/sdk/wandb_config.py

@@ -61,7 +61,7 @@ class Config:
 
         Using absl flags
         ```
-        flags.DEFINE_string(‘model’, None, ‘model to run’) # name, default, help
+        flags.DEFINE_string("model", None, "model to run") # name, default, help
         wandb.config.update(flags.FLAGS) # adds all absl flags to config
         ```