wandb 0.17.4__py3-none-win_amd64.whl → 0.17.6__py3-none-win_amd64.whl

wandb/sdk/launch/inputs/manage.py

@@ -1,12 +1,13 @@
 """Functions for declaring overridable configuration for launch jobs."""
 
-from typing import List, Optional
+from typing import Any, List, Optional
 
 
 def manage_config_file(
     path: str,
     include: Optional[List[str]] = None,
     exclude: Optional[List[str]] = None,
+    schema: Optional[Any] = None,
 ):
     r"""Declare an overridable configuration file for a launch job.
 
@@ -43,18 +44,27 @@ def manage_config_file(
             relative and must not contain backwards traversal, i.e. `..`.
         include (List[str]): A list of keys to include in the configuration file.
         exclude (List[str]): A list of keys to exclude from the configuration file.
+        schema (dict | Pydantic model): A JSON Schema or Pydantic model describing
+            describing which attributes will be editable from the Launch drawer.
+            Accepts both an instance of a Pydantic BaseModel class or the BaseModel
+            class itself.
 
     Raises:
         LaunchError: If the path is not valid, or if there is no active run.
     """
+    # note: schema's Any type is because in the case where a BaseModel class is
+    # provided, its type is a pydantic internal type that we don't want our typing
+    # to depend on. schema's type should be considered
+    # "Optional[dict | <something with a .model_json_schema() method>]"
     from .internal import handle_config_file_input
 
-    return handle_config_file_input(path, include, exclude)
+    return handle_config_file_input(path, include, exclude, schema)
 
 
 def manage_wandb_config(
     include: Optional[List[str]] = None,
     exclude: Optional[List[str]] = None,
+    schema: Optional[Any] = None,
 ):
     r"""Declare wandb.config as an overridable configuration for a launch job.
 
@@ -86,10 +96,18 @@ def manage_wandb_config(
     Args:
         include (List[str]): A list of subtrees to include in the configuration.
         exclude (List[str]): A list of subtrees to exclude from the configuration.
+        schema (dict | Pydantic model): A JSON Schema or Pydantic model describing
+            describing which attributes will be editable from the Launch drawer.
+            Accepts both an instance of a Pydantic BaseModel class or the BaseModel
+            class itself.
 
     Raises:
         LaunchError: If there is no active run.
     """
+    # note: schema's Any type is because in the case where a BaseModel class is
+    # provided, its type is a pydantic internal type that we don't want our typing
+    # to depend on. schema's type should be considered
+    # "Optional[dict | <something with a .model_json_schema() method>]"
     from .internal import handle_run_config_input
 
-    handle_run_config_input(include, exclude)
+    handle_run_config_input(include, exclude, schema)

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/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/lib/tracelog.py

@@ -45,8 +45,8 @@ logger = logging.getLogger(__name__)
 ANNOTATE_QUEUE_NAME = "_DEBUGLOG_QUEUE_NAME"
 
 # capture stdout and stderr before anyone messes with them
-stdout_write = sys.__stdout__.write
-stderr_write = sys.__stderr__.write
+stdout_write = sys.__stdout__.write  # type: ignore
+stderr_write = sys.__stderr__.write  # type: ignore
 
 
 def _log(

wandb/sdk/wandb_init.py

@@ -30,15 +30,7 @@ from wandb.util import _is_artifact_representation
 
 from . import wandb_login, wandb_setup
 from .backend.backend import Backend
-from .lib import (
-    RunDisabled,
-    SummaryDisabled,
-    filesystem,
-    ipython,
-    module,
-    reporting,
-    telemetry,
-)
+from .lib import SummaryDisabled, filesystem, ipython, module, reporting, telemetry
 from .lib.deprecate import Deprecated, deprecate
 from .lib.mailbox import Mailbox, MailboxProgress
 from .lib.printer import Printer, get_printer
@@ -529,20 +521,18 @@ class _WandbInit:
         logger.info(f"Logging user logs to {settings.log_user}")
         logger.info(f"Logging internal logs to {settings.log_internal}")
 
-    def _make_run_disabled(self) -> RunDisabled:
-        drun = RunDisabled()
-        drun.config = wandb.wandb_sdk.wandb_config.Config()
-        drun.config.update(self.sweep_config)
-        drun.config.update(self.config)
-        drun.summary = SummaryDisabled()
-        drun.log = lambda data, *_, **__: drun.summary.update(data)
-        drun.finish = lambda *_, **__: module.unset_globals()
-        drun.step = 0
-        drun.resumed = False
-        drun.disabled = True
-        drun.id = runid.generate_id()
-        drun.name = "dummy-" + drun.id
-        drun.dir = tempfile.gettempdir()
+    def _make_run_disabled(self) -> Run:
+        drun = Run(settings=Settings(mode="disabled", files_dir=tempfile.gettempdir()))
+        drun._config = wandb.wandb_sdk.wandb_config.Config()
+        drun._config.update(self.sweep_config)
+        drun._config.update(self.config)
+        drun.summary = SummaryDisabled()  # type: ignore
+        drun.log = lambda data, *_, **__: drun.summary.update(data)  # type: ignore
+        drun.finish = lambda *_, **__: module.unset_globals()  # type: ignore
+        drun._step = 0
+        drun._run_obj = None
+        drun._run_id = runid.generate_id()
+        drun._name = "dummy-" + drun.id
         module.set_global(
             run=drun,
             config=drun.config,
@@ -563,7 +553,7 @@ class _WandbInit:
         percent_done = handle.percent_done
         self.printer.progress_update(line, percent_done=percent_done)
 
-    def init(self) -> Union[Run, RunDisabled]:  # noqa: C901
+    def init(self) -> Run:  # noqa: C901
         if logger is None:
             raise RuntimeError("Logger not initialized")
         logger.info("calling init triggers")
@@ -700,6 +690,12 @@ class _WandbInit:
                 tel.feature.flow_control_custom = True
             if self.settings._require_core:
                 tel.feature.core = True
+            if self.settings._shared:
+                wandb.termwarn(
+                    "The `_shared` feature is experimental and may change. "
+                    "Please contact support@wandb.com for guidance and to report any issues."
+                )
+                tel.feature.shared_mode = True
 
             tel.env.maybe_mp = _maybe_mp_process(backend)
 
@@ -853,7 +849,7 @@ def _attach(
     run_id: Optional[str] = None,
     *,
     run: Optional["Run"] = None,
-) -> Union[Run, RunDisabled, None]:
+) -> Optional[Run]:
     """Attach to a run currently executing in another process/thread.
 
     Arguments:
@@ -907,7 +903,7 @@ def _attach(
     if run is None:
         run = Run(settings=settings)
     else:
-        run._init(settings=settings)
+        run._init()
     run._set_library(_wl)
     run._set_backend(backend)
     backend._hack_set_run(run)
@@ -957,7 +953,7 @@ def init(
     fork_from: Optional[str] = None,
     resume_from: Optional[str] = None,
     settings: Union[Settings, Dict[str, Any], None] = None,
-) -> Union[Run, RunDisabled]:
+) -> Run:
     r"""Start a new run to track and log to W&B.
 
     In an ML training pipeline, you could add `wandb.init()`

wandb/sdk/wandb_login.py

@@ -54,16 +54,16 @@ def login(
 ) -> bool:
     """Set up W&B login credentials.
 
-    By default, this will only store the credentials locally without
+    By default, this will only store credentials locally without
     verifying them with the W&B server. To verify credentials, pass
-    verify=True.
+    `verify=True`.
 
     Arguments:
         anonymous: (string, optional) Can be "must", "allow", or "never".
-            If set to "must" we'll always log in anonymously, if set to
-            "allow" we'll only create an anonymous user if the user
-            isn't already logged in.
-        key: (string, optional) authentication key.
+            If set to "must", always log a user in anonymously. If set to
+            "allow", only create an anonymous user if the user
+            isn't already logged in. If set to "never", never log a
+            user anonymously. Default set to "never".
         relogin: (bool, optional) If true, will re-prompt for API key.
         host: (string, optional) The host to connect to.
         force: (bool, optional) If true, will force a relogin.

wandb/sdk/wandb_manager.py

@@ -114,17 +114,21 @@ class _Manager:
 
         try:
             svc_iface._svc_connect(port=port)
+
         except ConnectionRefusedError as e:
             if not psutil.pid_exists(self._token.pid):
                 message = (
-                    "Connection to wandb service failed "
-                    "since the process is not available. "
+                    "Connection to wandb service failed"
+                    " because the process is not available."
                 )
             else:
-                message = f"Connection to wandb service failed: {e}. "
-            raise ManagerConnectionRefusedError(message)
+                message = "Connection to wandb service failed."
+            raise ManagerConnectionRefusedError(message) from e
+
         except Exception as e:
-            raise ManagerConnectionError(f"Connection to wandb service failed: {e}")
+            raise ManagerConnectionError(
+                "Connection to wandb service failed.",
+            ) from e
 
     def __init__(self, settings: "Settings") -> None:
         """Connects to the internal service, starting it if necessary."""