wandb 0.17.6__py3-none-win32.whl → 0.17.7__py3-none-win32.whl

wandb/__init__.py

@@ -1,9 +1,6 @@
 """Use wandb to track machine learning work.
 
-The most commonly used functions/objects are:
-  - wandb.init — initialize a new run at the top of your training script
-  - wandb.config — track hyperparameters and metadata
-  - wandb.log — log metrics and media over time within your training loop
+Train and fine-tune models, manage models from experimentation to production.
 
 For guides and examples, see https://docs.wandb.ai.
 
@@ -11,11 +8,7 @@ For scripts and interactive notebooks, see https://github.com/wandb/examples.
 
 For reference documentation, see https://docs.wandb.com/ref/python.
 """
-__version__ = "0.17.6"
-
-
-# Used with pypi checks and other messages related to pip
-_wandb_module = "wandb"
+__version__ = "0.17.7"
 
 from typing import Optional
 
@@ -118,13 +111,6 @@ def _assert_is_user_process():
     assert not _IS_INTERNAL_PROCESS
 
 
-# toplevel:
-# save()
-# restore()
-# login()
-# sweep()
-# agent()
-
 # globals
 Api = PublicApi
 api = InternalApi()
@@ -254,4 +240,5 @@ __all__ = (
     "log_model",
     "use_model",
     "link_model",
+    "define_metric",
 )

wandb/agents/pyagent.py

@@ -14,7 +14,6 @@ import time
 import traceback
 
 import wandb
-from wandb import wandb_sdk
 from wandb.apis import InternalApi
 from wandb.sdk.launch.sweeps import utils as sweep_utils
 
@@ -298,7 +297,7 @@ class Agent:
                 sweep_param_path, job.config
             )
             os.environ[wandb.env.SWEEP_ID] = self._sweep_id
-            wandb_sdk.wandb_setup._setup(_reset=True)
+            wandb.sdk.wandb_setup._setup(_reset=True)
 
             wandb.termlog(f"Agent Starting Run: {run_id} with config:")
             for k, v in job.config.items():

wandb/cli/cli.py

@@ -1417,6 +1417,13 @@ def launch_sweep(
     help="""Specific docker image you'd like to use. In the form name:tag.
     If passed in, will override the docker image value passed in using a config file.""",
 )
+@click.option(
+    "--base-image",
+    "-B",
+    default=None,
+    metavar="BASE IMAGE",
+    help="""Docker image to run job code in. Incompatible with --docker-image.""",
+)
 @click.option(
     "--config",
     "-c",
@@ -1508,6 +1515,7 @@ def launch(
     entity,
     project,
     docker_image,
+    base_image,
     config,
     cli_template_vars,
     queue,
@@ -1604,6 +1612,7 @@ def launch(
             git_hash=git_version,
             name=job_name,
             project=project,
+            base_image=base_image,
             build_context=build_context,
             dockerfile=dockerfile,
             entity=entity,
@@ -2001,6 +2010,12 @@ def describe(job):
     help="Path to the build context from the root of the job source code. If "
     "provided, this is used as the base path for the Dockerfile and entrypoint.",
 )
+@click.option(
+    "--base-image",
+    "-B",
+    type=str,
+    help="Base image to use for the job. Incompatible with image jobs.",
+)
 @click.option(
     "--dockerfile",
     "-D",
@@ -2025,6 +2040,7 @@ def create(
     git_hash,
     runtime,
     build_context,
+    base_image,
     dockerfile,
 ):
     """Create a job from a source, without a wandb run.
@@ -2056,6 +2072,10 @@ def create(
         )
         entrypoint = "main.py"
 
+    if job_type == "image" and base_image:
+        wandb.termerror("Cannot provide --base-image/-B for an `image` job")
+        return
+
     artifact, action, aliases = _create_job(
         api=api,
         path=path,
@@ -2069,6 +2089,7 @@ def create(
         git_hash=git_hash,
         runtime=runtime,
         build_context=build_context,
+        base_image=base_image,
         dockerfile=dockerfile,
     )
     if not artifact:

wandb/data_types.py

@@ -104,7 +104,7 @@ class _TableIndex(int, _TableLinkMixin):
 def _json_helper(val, artifact):
     if isinstance(val, WBValue):
         return val.to_json(artifact)
-    elif val.__class__ == dict:
+    elif val.__class__ is dict:
         res = {}
         for key in val:
             res[key] = _json_helper(val[key], artifact)
@@ -270,10 +270,10 @@ class Table(Media):
         if dtype is None:
             dtype = _dtypes.UnknownType()
 
-        if optional.__class__ != list:
+        if optional.__class__ is not list:
             optional = [optional for _ in range(len(self.columns))]
 
-        if dtype.__class__ != list:
+        if dtype.__class__ is not list:
             dtype = [dtype for _ in range(len(self.columns))]
 
         self._column_types = _dtypes.TypedDictType({})

wandb/integration/kfp/wandb_logging.py

@@ -30,7 +30,7 @@ def wandb_log(  # noqa: C901
     def isinstance_namedtuple(x):
         t = type(x)
         b = t.__bases__
-        if len(b) != 1 or b[0] != tuple:
+        if len(b) != 1 or b[0] is not tuple:
             return False
         f = getattr(t, "_fields", None)
         if not isinstance(f, tuple):

wandb/integration/lightning/fabric/logger.py

@@ -312,7 +312,7 @@ class WandbLogger(Logger):
         self._prefix = prefix
         self._experiment = experiment
         self._logged_model_time: Dict[str, float] = {}
-        self._checkpoint_callback: Optional["ModelCheckpoint"] = None
+        self._checkpoint_callback: Optional[ModelCheckpoint] = None
 
         # paths are processed as strings
         if save_dir is not None:

wandb/integration/openai/fine_tuning.py

@@ -1,3 +1,4 @@
+import base64
 import datetime
 import io
 import json
@@ -16,13 +17,13 @@ from wandb.util import parse_version
 
 openai = util.get_module(
     name="openai",
-    required="`openai` not installed. This integration requires `openai`. To fix, please `pip install openai`",
-    lazy="False",
+    required="This integration requires `openai`. To install, please run `pip install openai`",
+    lazy=False,
 )
 
-if parse_version(openai.__version__) < parse_version("1.0.1"):
+if parse_version(openai.__version__) < parse_version("1.12.0"):
     raise wandb.Error(
-        f"This integration requires openai version 1.0.1 and above. Your current version is {openai.__version__} "
+        f"This integration requires openai version 1.12.0 and above. Your current version is {openai.__version__} "
         "To fix, please `pip install -U openai`"
     )
 
@@ -228,7 +229,14 @@ class WandbLogger:
         # check results are present
         try:
             results_id = fine_tune.result_files[0]
-            results = cls.openai_client.files.content(file_id=results_id).text
+            try:
+                encoded_results = cls.openai_client.files.content(
+                    file_id=results_id
+                ).read()
+                results = base64.b64decode(encoded_results).decode("utf-8")
+            except Exception:
+                # attempt to read as text, works for older jobs
+                results = cls.openai_client.files.content(file_id=results_id).text
         except openai.NotFoundError:
             if show_individual_warnings:
                 wandb.termwarn(

wandb/integration/ultralytics/pose_utils.py

@@ -23,7 +23,6 @@ def annotate_keypoint_results(result: Results, visualize_skeleton: bool):
 
 
 def annotate_keypoint_batch(image_path: str, keypoints: Any, visualize_skeleton: bool):
-    original_image = None
     with Image.open(image_path) as original_image:
         original_image = np.ascontiguousarray(original_image)
         annotator = Annotator(original_image)

wandb/sdk/artifacts/artifact.py

@@ -751,7 +751,7 @@ class Artifact:
     def save(
         self,
         project: Optional[str] = None,
-        settings: Optional["wandb.wandb_sdk.wandb_settings.Settings"] = None,
+        settings: Optional["wandb.sdk.wandb_settings.Settings"] = None,
     ) -> None:
         """Persist any changes made to the artifact.
 

wandb/sdk/data_types/_dtypes.py

@@ -56,7 +56,7 @@ class TypeRegistry:
         # but will be ultimately treated as a None. Ignoring type since
         # mypy does not trust that py_obj is a float by the time it is
         # passed to isnan.
-        if py_obj.__class__ == float and math.isnan(py_obj):  # type: ignore
+        if py_obj.__class__ is float and math.isnan(py_obj):  # type: ignore
             return NoneType()
 
         # TODO: generalize this to handle other config input types
@@ -134,7 +134,7 @@ def _params_obj_to_json_obj(
     artifact: t.Optional["Artifact"] = None,
 ) -> t.Any:
     """Helper method."""
-    if params_obj.__class__ == dict:
+    if params_obj.__class__ is dict:
         return {
             key: _params_obj_to_json_obj(params_obj[key], artifact)
             for key in params_obj
@@ -151,7 +151,7 @@ def _json_obj_to_params_obj(
     json_obj: t.Any, artifact: t.Optional["Artifact"] = None
 ) -> t.Any:
     """Helper method."""
-    if json_obj.__class__ == dict:
+    if json_obj.__class__ is dict:
         if "wb_type" in json_obj:
             return TypeRegistry.type_from_dict(json_obj, artifact)
         else:
@@ -159,7 +159,7 @@ def _json_obj_to_params_obj(
                 key: _json_obj_to_params_obj(json_obj[key], artifact)
                 for key in json_obj
             }
-    elif json_obj.__class__ == list:
+    elif json_obj.__class__ is list:
         return [_json_obj_to_params_obj(item, artifact) for item in json_obj]
     else:
         return json_obj
@@ -533,7 +533,7 @@ class UnionType(Type):
         self,
         allowed_types: t.Optional[t.Sequence[ConvertibleToType]] = None,
     ):
-        assert allowed_types is None or (allowed_types.__class__ == list)
+        assert allowed_types is None or (allowed_types.__class__ is list)
         if allowed_types is None:
             wb_types = []
         else:

wandb/sdk/data_types/base_types/media.py

@@ -159,9 +159,11 @@ class Media(WBValue):
         # into Media itself we should get rid of them
         from wandb import Image
         from wandb.data_types import Audio
+        from wandb.sdk.wandb_run import Run
 
         json_obj = {}
-        if isinstance(run, wandb.wandb_sdk.wandb_run.Run):
+
+        if isinstance(run, Run):
             json_obj.update(
                 {
                     "_type": "file",  # TODO(adrian): This isn't (yet) a real media type we support on the frontend.

wandb/sdk/data_types/helper_types/bounding_boxes_2d.py

@@ -276,7 +276,9 @@ class BoundingBoxes2D(JSONMetadata):
         return True
 
     def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
-        if isinstance(run_or_artifact, wandb.wandb_sdk.wandb_run.Run):
+        from wandb.sdk.wandb_run import Run
+
+        if isinstance(run_or_artifact, Run):
             return super().to_json(run_or_artifact)
         elif isinstance(run_or_artifact, wandb.Artifact):
             # TODO (tim): I would like to log out a proper dictionary representing this object, but don't

wandb/sdk/data_types/helper_types/image_mask.py

@@ -191,9 +191,11 @@ class ImageMask(Media):
         )
 
     def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
+        from wandb.sdk.wandb_run import Run
+
         json_dict = super().to_json(run_or_artifact)
 
-        if isinstance(run_or_artifact, wandb.wandb_sdk.wandb_run.Run):
+        if isinstance(run_or_artifact, Run):
             json_dict["_type"] = self.type_name()
             return json_dict
         elif isinstance(run_or_artifact, wandb.Artifact):

wandb/sdk/data_types/image.py

@@ -418,6 +418,8 @@ class Image(BatchableMedia):
                 )
 
     def to_json(self, run_or_artifact: Union["LocalRun", "Artifact"]) -> dict:
+        from wandb.sdk.wandb_run import Run
+
         json_dict = super().to_json(run_or_artifact)
         json_dict["_type"] = Image._log_type
         json_dict["format"] = self.format
@@ -456,7 +458,7 @@ class Image(BatchableMedia):
                     "digest": classes_entry.digest,
                 }
 
-        elif not isinstance(run_or_artifact, wandb.wandb_sdk.wandb_run.Run):
+        elif not isinstance(run_or_artifact, Run):
             raise ValueError("to_json accepts wandb_run.Run or wandb_artifact.Artifact")
 
         if self._boxes:

wandb/sdk/data_types/saved_model.py

@@ -148,7 +148,9 @@ class _SavedModel(WBValue, Generic[SavedModelObjType]):
         # bit of tech debt in the other data types which requires the input to `to_json`
         # to accept a Run or Artifact. However, Run additions should be deprecated in the future.
         # This check helps ensure we do not add to the debt.
-        if isinstance(run_or_artifact, wandb.wandb_sdk.wandb_run.Run):
+        from wandb.sdk.wandb_run import Run
+
+        if isinstance(run_or_artifact, Run):
             raise ValueError("SavedModel cannot be added to run - must use artifact")
         artifact = run_or_artifact
         json_obj = {

wandb/sdk/interface/interface.py

@@ -54,18 +54,20 @@ MANIFEST_FILE_SIZE_THRESHOLD = 100_000
 
 GlobStr = NewType("GlobStr", str)
 
-if TYPE_CHECKING:
-    from ..wandb_run import Run
+if sys.version_info >= (3, 8):
+    from typing import Literal, TypedDict
+else:
+    from typing_extensions import Literal, TypedDict
+
+PolicyName = Literal["now", "live", "end"]
 
-    if sys.version_info >= (3, 8):
-        from typing import Literal, TypedDict
-    else:
-        from typing_extensions import Literal, TypedDict
 
-    PolicyName = Literal["now", "live", "end"]
+class FilesDict(TypedDict):
+    files: Iterable[Tuple[GlobStr, PolicyName]]
 
-    class FilesDict(TypedDict):
-        files: Iterable[Tuple[GlobStr, PolicyName]]
+
+if TYPE_CHECKING:
+    from ..wandb_run import Run
 
 
 logger = logging.getLogger("wandb")
@@ -112,15 +114,14 @@ class InterfaceBase:
     def _publish_header(self, header: pb.HeaderRecord) -> None:
         raise NotImplementedError
 
-    def communicate_status(self) -> Optional[pb.StatusResponse]:
-        status = pb.StatusRequest()
-        resp = self._communicate_status(status)
-        return resp
+    def deliver_status(self) -> MailboxHandle:
+        return self._deliver_status(pb.StatusRequest())
 
     @abstractmethod
-    def _communicate_status(
-        self, status: pb.StatusRequest
-    ) -> Optional[pb.StatusResponse]:
+    def _deliver_status(
+        self,
+        status: pb.StatusRequest,
+    ) -> MailboxHandle:
         raise NotImplementedError
 
     def _make_config(

wandb/sdk/interface/interface_shared.py

@@ -299,7 +299,7 @@ class InterfaceShared(InterfaceBase):
         raise NotImplementedError
 
     def _communicate(
-        self, rec: pb.Record, timeout: Optional[int] = 5, local: Optional[bool] = None
+        self, rec: pb.Record, timeout: Optional[int] = 30, local: Optional[bool] = None
     ) -> Optional[pb.Result]:
         return self._communicate_async(rec, local=local).get(timeout=timeout)
 
@@ -421,15 +421,12 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(alert=proto_alert)
         self._publish(rec)
 
-    def _communicate_status(
-        self, status: pb.StatusRequest
-    ) -> Optional[pb.StatusResponse]:
+    def _deliver_status(
+        self,
+        status: pb.StatusRequest,
+    ) -> MailboxHandle:
         req = self._make_request(status=status)
-        resp = self._communicate(req, local=True)
-        if resp is None:
-            return None
-        assert resp.response.status_response
-        return resp.response.status_response
+        return self._deliver_record(req)
 
     def _publish_exit(self, exit_data: pb.RunExitRecord) -> None:
         rec = self._make_record(exit=exit_data)

wandb/sdk/internal/datastore.py

@@ -69,7 +69,7 @@ class DataStore:
 
     def __init__(self) -> None:
         self._opened_for_scan = False
-        self._fp: Optional["IO[Any]"] = None
+        self._fp: Optional[IO[Any]] = None
         self._index = 0
         self._flush_offset = 0
         self._size_bytes = 0

wandb/sdk/internal/handler.py

@@ -745,8 +745,6 @@ class HandleManager:
         self._respond_result(result)
 
     def handle_request_status(self, record: Record) -> None:
-        # TODO(mempressure): do something better?
-        assert record.control.req_resp
         result = proto_util._result_from_record(record)
         self._respond_result(result)
 

wandb/sdk/internal/internal.py

@@ -62,7 +62,7 @@ def wandb_internal(
 
     """
     # mark this process as internal
-    wandb._set_internal_process()
+    wandb._set_internal_process()  # type: ignore
     _setup_tracelog()
     started = time.time()
 

wandb/sdk/internal/job_builder.py

@@ -423,15 +423,18 @@ class JobBuilder:
         api: Api,
         build_context: Optional[str] = None,
         dockerfile: Optional[str] = None,
+        base_image: Optional[str] = None,
     ) -> Optional[Artifact]:
         """Build a job artifact from the current run.
 
         Arguments:
+            api (Api): The API object to use to create the job artifact.
             build_context (Optional[str]): Path within the job source code to
                 the image build context. Saved as part of the job for future
                 builds.
             dockerfile (Optional[str]): Path within the build context the
                 Dockerfile. Saved as part of the job for future builds.
+            base_image (Optional[str]): The base image used to run the job code.
 
         Returns:
             Optional[Artifact]: The job artifact if it was successfully built,
@@ -467,8 +470,6 @@ class JobBuilder:
                 "warn",
             )
             return None
-        metadata["dockerfile"] = dockerfile
-        metadata["build_context"] = build_context
 
         runtime: Optional[str] = metadata.get("python")
         # can't build a job without a python version
@@ -520,6 +521,8 @@ class JobBuilder:
             source["build_context"] = build_context  # type: ignore[typeddict-item]
         if dockerfile:
             source["dockerfile"] = dockerfile  # type: ignore[typeddict-item]
+        if base_image:
+            source["base_image"] = base_image  # type: ignore[typeddict-item]
 
         # Pop any keys that are initialized to None. The current TypedDict
         # system for source dicts requires all keys to be present, but we

wandb/sdk/internal/tb_watcher.py

@@ -123,7 +123,7 @@ class TBWatcher:
         self._force = force
         # TODO(jhr): do we need locking in this queue?
         self._watcher_queue = queue.PriorityQueue()
-        wandb.tensorboard.reset_state()
+        wandb.tensorboard.reset_state()  # type: ignore
 
     def _calculate_namespace(self, logdir: str, rootdir: str) -> Optional[str]:
         namespace: Optional[str]
@@ -430,7 +430,7 @@ class TBEventConsumer:
     def _handle_event(
         self, event: "ProtoEvent", history: Optional["TBHistory"] = None
     ) -> None:
-        wandb.tensorboard._log(
+        wandb.tensorboard._log(  # type: ignore
             event.event,
             step=event.event.step,
             namespace=event.namespace,

wandb/sdk/internal/update.py

@@ -10,7 +10,7 @@ def _find_available(
 ) -> Optional[Tuple[str, bool, bool, bool, Optional[str]]]:
     from wandb.util import parse_version
 
-    pypi_url = f"https://pypi.org/pypi/{wandb._wandb_module}/json"
+    pypi_url = "https://pypi.org/pypi/wandb/json"
 
     yanked_dict = {}
     try:
@@ -78,7 +78,7 @@ def check_available(current_version: str) -> Optional[Dict[str, Optional[str]]]:
     if not package_info:
         return None
 
-    wandb_module_name = wandb._wandb_module
+    wandb_module_name = "wandb"
 
     latest_version, pip_prerelease, deleted, yanked, yanked_reason = package_info
     upgrade_message = (

wandb/sdk/launch/builder/kaniko_builder.py

@@ -63,6 +63,13 @@ else:
     NAMESPACE = "wandb"
 
 
+def get_pod_name_safe(job: client.V1Job):
+    try:
+        return job.spec.template.metadata.name
+    except AttributeError:
+        return None
+
+
 async def _wait_for_completion(
     batch_client: client.BatchV1Api, job_name: str, deadline_secs: Optional[int] = None
 ) -> bool:
@@ -319,17 +326,18 @@ class KanikoBuilder(AbstractBuilder):
                 await self._create_docker_ecr_config_map(
                     build_job_name, core_v1, repo_uri
                 )
-            await batch_v1.create_namespaced_job(NAMESPACE, build_job)
-
+            k8s_job = await batch_v1.create_namespaced_job(NAMESPACE, build_job)
             # wait for double the job deadline since it might take time to schedule
             if not await _wait_for_completion(
                 batch_v1, build_job_name, 3 * _DEFAULT_BUILD_TIMEOUT_SECS
             ):
                 if job_tracker:
                     job_tracker.set_err_stage("build")
-                raise Exception(
-                    f"Failed to build image in kaniko for job {run_id}. View logs with `kubectl logs -n {NAMESPACE} {build_job_name}`."
-                )
+                msg = f"Failed to build image in kaniko for job {run_id}."
+                pod_name = get_pod_name_safe(k8s_job)
+                if pod_name:
+                    msg += f" View logs with `kubectl logs -n {NAMESPACE} {pod_name}`."
+                raise Exception(msg)
             try:
                 pods_from_job = await core_v1.list_namespaced_pod(
                     namespace=NAMESPACE, label_selector=f"job-name={build_job_name}"

wandb/sdk/launch/create_job.py

@@ -114,6 +114,7 @@ def _create_job(
     git_hash: Optional[str] = None,
     build_context: Optional[str] = None,
     dockerfile: Optional[str] = None,
+    base_image: Optional[str] = None,
 ) -> Tuple[Optional[Artifact], str, List[str]]:
     wandb.termlog(f"Creating launch job of type: {job_type}...")
 
@@ -188,6 +189,7 @@ def _create_job(
         api.api,
         dockerfile=dockerfile,
         build_context=build_context,
+        base_image=base_image,
     )
     if not artifact:
         wandb.termerror("JobBuilder failed to build a job")

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/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
         ```