wandb 0.16.4__py3-none-any.whl → 0.16.6__py3-none-any.whl

wandb/sdk/interface/interface.py

@@ -13,8 +13,19 @@ import os
 import sys
 import time
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, Iterable, NewType, Optional, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterable,
+    List,
+    NewType,
+    Optional,
+    Tuple,
+    Union,
+)
 
+from wandb import termwarn
 from wandb.proto import wandb_internal_pb2 as pb
 from wandb.proto import wandb_telemetry_pb2 as tpb
 from wandb.sdk.artifacts.artifact import Artifact
@@ -340,6 +351,7 @@ class InterfaceBase:
                 proto_entry.ref = entry.ref
             if entry.local_path:
                 proto_entry.local_path = entry.local_path
+            proto_entry.skip_cache = entry.skip_cache
             for k, v in entry.extra.items():
                 proto_extra = proto_entry.extra.add()
                 proto_extra.key = k
@@ -436,16 +448,27 @@ class InterfaceBase:
                 path = artifact.get_entry("wandb-job.json").download()
                 with open(path) as f:
                     job_info = json.load(f)
+
             except Exception as e:
                 logger.warning(
                     f"Failed to download partial job info from artifact {artifact}, : {e}"
                 )
-            use_artifact = self._make_proto_use_artifact(
-                use_artifact=use_artifact,
-                job_name=artifact.name,
-                job_info=job_info,
-                metadata=artifact.metadata,
-            )
+                termwarn(
+                    f"Failed to download partial job info from artifact {artifact}, : {e}"
+                )
+                return
+
+            try:
+                use_artifact = self._make_proto_use_artifact(
+                    use_artifact=use_artifact,
+                    job_name=artifact.name,
+                    job_info=job_info,
+                    metadata=artifact.metadata,
+                )
+            except Exception as e:
+                logger.warning(f"Failed to construct use artifact proto: {e}")
+                termwarn(f"Failed to construct use artifact proto: {e}")
+                return
 
         self._publish_use_artifact(use_artifact)
 
@@ -756,6 +779,36 @@ class InterfaceBase:
         run_start.run.CopyFrom(run_pb)
         return self._deliver_run_start(run_start)
 
+    def publish_launch_wandb_config_parameters(
+        self, include_paths: List[List[str]], exclude_paths: List[List[str]]
+    ):
+        """Tells the internal process to treat wandb.config fields as job inputs.
+
+        The paths provided as arguments are sequences of dictionary keys that
+        specify a path within the wandb.config. If a path is included, the
+        corresponding field will be treated as a job input. If a path is
+        excluded, the corresponding field will not be treated as a job input.
+
+        Args:
+            include_paths: paths within config to include as job inputs.
+            exclude_paths: paths within config to exclude as job inputs.
+
+        Returns:
+            None
+        """
+        config_parameters = pb.LaunchWandbConfigParametersRecord()
+        include_records = [pb.ConfigFilterPath(path=path) for path in include_paths]
+        exclude_records = [pb.ConfigFilterPath(path=path) for path in exclude_paths]
+        config_parameters.include_paths.extend(include_records)
+        config_parameters.exclude_paths.extend(exclude_records)
+        return self._publish_launch_wandb_config_parameters(config_parameters)
+
+    @abstractmethod
+    def _publish_launch_wandb_config_parameters(
+        self, config_parameters: pb.LaunchWandbConfigParametersRecord
+    ) -> None:
+        raise NotImplementedError
+
     @abstractmethod
     def _deliver_run_start(self, run_start: pb.RunStartRequest) -> MailboxHandle:
         raise NotImplementedError
@@ -871,11 +924,3 @@ class InterfaceBase:
         self, run_status: pb.RunStatusRequest
     ) -> MailboxHandle:
         raise NotImplementedError
-
-    def deliver_request_job_info(self) -> MailboxHandle:
-        job_info = pb.JobInfoRequest()
-        return self._deliver_request_job_info(job_info)
-
-    @abstractmethod
-    def _deliver_request_job_info(self, job_info: pb.JobInfoRequest) -> MailboxHandle:
-        raise NotImplementedError

wandb/sdk/interface/interface_shared.py

@@ -145,7 +145,6 @@ class InterfaceShared(InterfaceBase):
         cancel: Optional[pb.CancelRequest] = None,
         summary_record: Optional[pb.SummaryRecordRequest] = None,
         telemetry_record: Optional[pb.TelemetryRecordRequest] = None,
-        job_info: Optional[pb.JobInfoRequest] = None,
         get_system_metrics: Optional[pb.GetSystemMetricsRequest] = None,
         python_packages: Optional[pb.PythonPackagesRequest] = None,
     ) -> pb.Record:
@@ -202,8 +201,6 @@ class InterfaceShared(InterfaceBase):
             request.summary_record.CopyFrom(summary_record)
         elif telemetry_record:
             request.telemetry_record.CopyFrom(telemetry_record)
-        elif job_info:
-            request.job_info.CopyFrom(job_info)
         elif get_system_metrics:
             request.get_system_metrics.CopyFrom(get_system_metrics)
         elif sync:
@@ -242,6 +239,9 @@ class InterfaceShared(InterfaceBase):
         use_artifact: Optional[pb.UseArtifactRecord] = None,
         output: Optional[pb.OutputRecord] = None,
         output_raw: Optional[pb.OutputRawRecord] = None,
+        launch_wandb_config_parameters: Optional[
+            pb.LaunchWandbConfigParametersRecord
+        ] = None,
     ) -> pb.Record:
         record = pb.Record()
         if run:
@@ -286,6 +286,8 @@ class InterfaceShared(InterfaceBase):
             record.output.CopyFrom(output)
         elif output_raw:
             record.output_raw.CopyFrom(output_raw)
+        elif launch_wandb_config_parameters:
+            record.wandb_config_parameters.CopyFrom(launch_wandb_config_parameters)
         else:
             raise Exception("Invalid record")
         return record
@@ -415,6 +417,14 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(alert=proto_alert)
         self._publish(rec)
 
+    def _publish_launch_wandb_config_parameters(
+        self, launch_wandb_config_parameters: pb.LaunchWandbConfigParametersRecord
+    ) -> None:
+        rec = self._make_record(
+            launch_wandb_config_parameters=launch_wandb_config_parameters
+        )
+        self._publish(rec)
+
     def _communicate_status(
         self, status: pb.StatusRequest
     ) -> Optional[pb.StatusResponse]:
@@ -523,10 +533,6 @@ class InterfaceShared(InterfaceBase):
         record = self._make_request(run_status=run_status)
         return self._deliver_record(record)
 
-    def _deliver_request_job_info(self, job_info: pb.JobInfoRequest) -> MailboxHandle:
-        record = self._make_request(job_info=job_info)
-        return self._deliver_record(record)
-
     def _transport_keepalive_failed(self, keepalive_interval: int = 5) -> bool:
         if self._transport_failed:
             return True

wandb/sdk/internal/file_stream.py

@@ -1,6 +1,7 @@
 import base64
 import functools
 import itertools
+import json
 import logging
 import os
 import queue
@@ -58,6 +59,7 @@ class Chunk(NamedTuple):
 class DefaultFilePolicy:
     def __init__(self, start_chunk_id: int = 0) -> None:
         self._chunk_id = start_chunk_id
+        self.has_debug_log = False
 
     def process_chunks(
         self, chunks: List[Chunk]
@@ -66,6 +68,21 @@ class DefaultFilePolicy:
         self._chunk_id += len(chunks)
         return {"offset": chunk_id, "content": [c.data for c in chunks]}
 
+    # TODO: this is very inefficient, this is meant for temporary debugging and will be removed in future releases
+    def _debug_log(self, data: Any):
+        if self.has_debug_log or not os.environ.get("WANDB_DEBUG_FILESTREAM_LOG"):
+            return
+
+        loaded = json.loads(data)
+        if not isinstance(loaded, dict):
+            return
+
+        # get key size and convert to MB
+        key_sizes = [(k, len(json.dumps(v))) for k, v in loaded.items()]
+        key_msg = [f"{k}: {v/1048576:.5f} MB" for k, v in key_sizes]
+        wandb.termerror(f"Step: {loaded['_step']} | {key_msg}", repeat=False)
+        self.has_debug_log = True
+
 
 class JsonlFilePolicy(DefaultFilePolicy):
     def process_chunks(self, chunks: List[Chunk]) -> "ProcessedChunk":
@@ -81,6 +98,7 @@ class JsonlFilePolicy(DefaultFilePolicy):
                 )
                 wandb.termerror(msg, repeat=False)
                 wandb._sentry.message(msg, repeat=False)
+                self._debug_log(chunk.data)
             else:
                 chunk_data.append(chunk.data)
 
@@ -99,6 +117,7 @@ class SummaryFilePolicy(DefaultFilePolicy):
             )
             wandb.termerror(msg, repeat=False)
             wandb._sentry.message(msg, repeat=False)
+            self._debug_log(data)
             return False
         return {"offset": 0, "content": [data]}
 

wandb/sdk/internal/handler.py

@@ -689,7 +689,7 @@ class HandleManager:
             self._settings, interface=self._interface, run_proto=run_start.run
         )
 
-        if run_start.run.resumed:
+        if run_start.run.resumed or run_start.run.forked:
             self._step = run_start.run.starting_step
         result = proto_util._result_from_record(record)
         self._respond_result(result)
@@ -862,9 +862,6 @@ class HandleManager:
         self._respond_result(result)
         self._stopped.set()
 
-    def handle_request_job_info(self, record: Record) -> None:
-        self._dispatch_record(record, always_send=True)
-
     def finish(self) -> None:
         logger.info("shutting down handler")
         if self._system_monitor is not None:

wandb/sdk/internal/internal_api.py

@@ -2150,6 +2150,7 @@ class Api:
                             name
                         }
                     }
+                    historyLineCount
                 }
                 inserted
                 _Server_Settings_
@@ -2237,6 +2238,7 @@ class Api:
                 .get("serverSettings", {})
                 .get("serverMessages", [])
             )
+
         return (
             response["upsertBucket"]["bucket"],
             response["upsertBucket"]["inserted"],

wandb/sdk/internal/job_builder.py

@@ -4,7 +4,7 @@ import logging
 import os
 import re
 import sys
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
 
 import wandb
 from wandb.sdk.artifacts.artifact import Artifact
@@ -28,6 +28,8 @@ FROZEN_REQUIREMENTS_FNAME = "requirements.frozen.txt"
 JOB_FNAME = "wandb-job.json"
 JOB_ARTIFACT_TYPE = "job"
 
+LOG_LEVEL = Literal["log", "warn", "error"]
+
 
 class GitInfo(TypedDict):
     remote: str
@@ -89,8 +91,9 @@ class JobBuilder:
     _job_seq_id: Optional[str]
     _job_version_alias: Optional[str]
     _is_notebook_run: bool
+    _verbose: bool
 
-    def __init__(self, settings: SettingsStatic):
+    def __init__(self, settings: SettingsStatic, verbose: bool = False):
         self._settings = settings
         self._metadatafile_path = None
         self._requirements_path = None
@@ -106,6 +109,7 @@ class JobBuilder:
             Literal["repo", "artifact", "image"]
         ] = settings.job_source  # type: ignore[assignment]
         self._is_notebook_run = self._get_is_notebook_run()
+        self._verbose = verbose
 
     def set_config(self, config: Dict[str, Any]) -> None:
         self._config = config
@@ -197,6 +201,21 @@ class JobBuilder:
 
         return source, name
 
+    def _log_if_verbose(self, message: str, level: LOG_LEVEL) -> None:
+        log_func: Optional[Union[Callable[[Any], None], Callable[[Any], None]]] = None
+        if level == "log":
+            _logger.info(message)
+            log_func = wandb.termlog
+        elif level == "warn":
+            _logger.warning(message)
+            log_func = wandb.termwarn
+        elif level == "error":
+            _logger.error(message)
+            log_func = wandb.termerror
+
+        if self._verbose and log_func is not None:
+            log_func(message)
+
     def _build_artifact_job_source(
         self,
         program_relpath: str,
@@ -212,8 +231,9 @@ class JobBuilder:
                 # at the directory the notebook is in instead of the jupyter core
                 if not os.path.exists(os.path.basename(program_relpath)):
                     _logger.info("target path does not exist, exiting")
-                    wandb.termwarn(
-                        "No program path found when generating artifact job source for a non-colab notebook run. See https://docs.wandb.ai/guides/launch/create-job"
+                    self._log_if_verbose(
+                        "No program path found when generating artifact job source for a non-colab notebook run. See https://docs.wandb.ai/guides/launch/create-job",
+                        "warn",
                     )
                     return None, None
                 full_program_relpath = os.path.basename(program_relpath)
@@ -299,22 +319,25 @@ class JobBuilder:
         if not os.path.exists(
             os.path.join(self._settings.files_dir, REQUIREMENTS_FNAME)
         ):
-            wandb.termwarn(
-                "No requirements.txt found, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job"
+            self._log_if_verbose(
+                "No requirements.txt found, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job",
+                "warn",
             )
             return None
         metadata = self._handle_metadata_file()
         if metadata is None:
-            wandb.termwarn(
-                f"Ensure read and write access to run files dir: {self._settings.files_dir}, control this via the WANDB_DIR env var. See https://docs.wandb.ai/guides/track/environment-variables"
+            self._log_if_verbose(
+                f"Ensure read and write access to run files dir: {self._settings.files_dir}, control this via the WANDB_DIR env var. See https://docs.wandb.ai/guides/track/environment-variables",
+                "warn",
             )
             return None
 
         runtime: Optional[str] = metadata.get("python")
         # can't build a job without a python version
         if runtime is None:
-            wandb.termwarn(
-                "No python version found in metadata, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job"
+            self._log_if_verbose(
+                "No python version found in metadata, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job",
+                "warn",
             )
             return None
 
@@ -345,13 +368,16 @@ class JobBuilder:
                     or self._settings.job_source
                     or self._source_type
                 ):
-                    wandb.termwarn("No source type found, not creating job artifact")
+                    self._log_if_verbose(
+                        "No source type found, not creating job artifact", "warn"
+                    )
                 return None
 
             program_relpath = self._get_program_relpath(source_type, metadata)
             if source_type != "image" and not program_relpath:
-                wandb.termwarn(
-                    "No program path found, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job"
+                self._log_if_verbose(
+                    "No program path found, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job",
+                    "warn",
                 )
                 return None
 
@@ -377,10 +403,11 @@ class JobBuilder:
 
             if source is None:
                 if source_type:
-                    wandb.termwarn(
+                    self._log_if_verbose(
                         f"Source type is set to '{source_type}' but some required information is missing "
                         "from the environment. A job will not be created from this run. See "
-                        "https://docs.wandb.ai/guides/launch/create-job"
+                        "https://docs.wandb.ai/guides/launch/create-job",
+                        "warn",
                     )
                 return None
 
@@ -447,8 +474,9 @@ class JobBuilder:
             program = metadata.get("program")
 
             if not program:
-                wandb.termwarn(
-                    "Notebook 'program' path not found in metadata. See https://docs.wandb.ai/guides/launch/create-job"
+                self._log_if_verbose(
+                    "Notebook 'program' path not found in metadata. See https://docs.wandb.ai/guides/launch/create-job",
+                    "warn",
                 )
 
             return program

wandb/sdk/internal/sender.py

@@ -115,6 +115,7 @@ def _manifest_json_from_proto(manifest: "ArtifactManifest") -> Dict:
                 "ref": content.ref if content.ref else None,
                 "size": content.size if content.size is not None else None,
                 "local_path": content.local_path if content.local_path else None,
+                "skip_cache": content.skip_cache,
                 "extra": {
                     extra.key: json.loads(extra.value_json) for extra in content.extra
                 },
@@ -733,18 +734,7 @@ class SendManager:
             )
         self._respond_result(result)
 
-    def send_request_job_info(self, record: "Record") -> None:
-        """Respond to a request for a job link."""
-        result = proto_util._result_from_record(record)
-        result.response.job_info_response.sequenceId = (
-            self._job_builder._job_seq_id or ""
-        )
-        result.response.job_info_response.version = (
-            self._job_builder._job_version_alias or ""
-        )
-        self._respond_result(result)
-
-    def _maybe_setup_resume(
+    def _setup_resume(
         self, run: "RunRecord"
     ) -> Optional["wandb_internal_pb2.ErrorInfo"]:
         """Queries the backend for a run; fail if the settings are incompatible."""
@@ -890,6 +880,30 @@ class SendManager:
             pass
         # TODO: do something if sync spell is not successful?
 
+    def _setup_fork(self, server_run: dict):
+        assert self._settings.fork_from
+        assert self._settings.fork_from.metric == "_step"
+        assert self._run
+        first_step = int(self._settings.fork_from.value) + 1
+        self._resume_state.step = first_step
+        self._resume_state.history = server_run.get("historyLineCount", 0)
+        self._run.forked = True
+        self._run.starting_step = first_step
+
+    def _handle_error(
+        self,
+        record: "Record",
+        error: "wandb_internal_pb2.ErrorInfo",
+        run: "RunRecord",
+    ) -> None:
+        if record.control.req_resp or record.control.mailbox_slot:
+            result = proto_util._result_from_record(record)
+            result.run_result.run.CopyFrom(run)
+            result.run_result.error.CopyFrom(error)
+            self._respond_result(result)
+        else:
+            logger.error("Got error in async mode: %s", error.message)
+
     def send_run(self, record: "Record", file_dir: Optional[str] = None) -> None:
         run = record.run
         error = None
@@ -911,21 +925,28 @@ class SendManager:
             config_value_dict = self._config_backend_dict()
             self._config_save(config_value_dict)
 
+        do_fork = self._settings.fork_from is not None and is_wandb_init
+        do_resume = bool(self._settings.resume)
+
+        if do_fork and do_resume:
+            error = wandb_internal_pb2.ErrorInfo()
+            error.code = wandb_internal_pb2.ErrorInfo.ErrorCode.USAGE
+            error.message = (
+                "You cannot use `resume` and `fork_from` together. Please choose one."
+            )
+            self._handle_error(record, error, run)
+
         if is_wandb_init:
             # Ensure we have a project to query for status
             if run.project == "":
                 run.project = util.auto_project_name(self._settings.program)
             # Only check resume status on `wandb.init`
-            error = self._maybe_setup_resume(run)
+
+            if do_resume:
+                error = self._setup_resume(run)
 
         if error is not None:
-            if record.control.req_resp or record.control.mailbox_slot:
-                result = proto_util._result_from_record(record)
-                result.run_result.run.CopyFrom(run)
-                result.run_result.error.CopyFrom(error)
-                self._respond_result(result)
-            else:
-                logger.error("Got error in async mode: %s", error.message)
+            self._handle_error(record, error, run)
             return
 
         # Save the resumed config
@@ -945,19 +966,22 @@ class SendManager:
             self._config_save(config_value_dict)
 
         try:
-            self._init_run(run, config_value_dict)
+            server_run = self._init_run(run, config_value_dict)
         except (CommError, UsageError) as e:
             logger.error(e, exc_info=True)
-            if record.control.req_resp or record.control.mailbox_slot:
-                result = proto_util._result_from_record(record)
-                result.run_result.run.CopyFrom(run)
-                error = ProtobufErrorHandler.from_exception(e)
-                result.run_result.error.CopyFrom(error)
-                self._respond_result(result)
+            error = ProtobufErrorHandler.from_exception(e)
+            self._handle_error(record, error, run)
             return
 
         assert self._run  # self._run is configured in _init_run()
 
+        if do_fork:
+            error = self._setup_fork(server_run)
+
+        if error is not None:
+            self._handle_error(record, error, run)
+            return
+
         if record.control.req_resp or record.control.mailbox_slot:
             result = proto_util._result_from_record(record)
             # TODO: we could do self._interface.publish_defer(resp) to notify
@@ -976,7 +1000,7 @@ class SendManager:
         self,
         run: "RunRecord",
         config_dict: Optional[sender_config.BackendConfigDict],
-    ) -> None:
+    ) -> dict:
         # We subtract the previous runs runtime when resuming
         start_time = (
             run.start_time.ToMicroseconds() / 1e6
@@ -1061,6 +1085,7 @@ class SendManager:
             self._run.sweep_id = sweep_id
         if os.getenv("SPELL_RUN_URL"):
             self._sync_spell()
+        return server_run
 
     def _start_run_threads(self, file_dir: Optional[str] = None) -> None:
         assert self._run  # self._run is configured by caller

wandb/sdk/internal/settings_static.py

@@ -2,6 +2,7 @@ from dataclasses import fields
 from typing import Any, Iterable, Sequence, Tuple
 
 from wandb.proto import wandb_settings_pb2
+from wandb.sdk.lib import RunMoment
 from wandb.sdk.wandb_settings import SettingsData
 
 
@@ -38,6 +39,14 @@ class SettingsStatic(SettingsData):
                             unpacked_inner[inner_key] = inner_value
                         unpacked_mapping[outer_key] = unpacked_inner
                     value = unpacked_mapping
+            elif key == "fork_from":
+                value = getattr(proto, key)
+                if value.run:
+                    value = RunMoment(
+                        run=value.run, value=value.value, metric=value.metric
+                    )
+                else:
+                    value = None
             else:
                 if proto.HasField(key):  # type: ignore [arg-type]
                     value = getattr(proto, key).value

wandb/sdk/internal/system/system_info.py

@@ -212,7 +212,10 @@ class SystemInfo:
                 os.path.join(self.settings.files_dir, CONDA_ENVIRONMENTS_FNAME), "w"
             ) as f:
                 subprocess.call(
-                    ["conda", "env", "export"], stdout=f, stderr=subprocess.DEVNULL
+                    ["conda", "env", "export"],
+                    stdout=f,
+                    stderr=subprocess.DEVNULL,
+                    timeout=15,  # add timeout since conda env export could take a really long time
                 )
         except Exception as e:
             logger.exception(f"Error saving conda packages: {e}")

wandb/sdk/launch/_launch.py

@@ -62,6 +62,7 @@ def resolve_agent_config(  # noqa: C901
     max_jobs: Optional[int],
     queues: Optional[Tuple[str]],
     config: Optional[str],
+    verbosity: Optional[int],
 ) -> Tuple[Dict[str, Any], Api]:
     """Resolve the agent config.
 
@@ -72,6 +73,7 @@ def resolve_agent_config(  # noqa: C901
         max_jobs (int): The max number of jobs.
         queues (Tuple[str]): The queues.
         config (str): The config.
+        verbosity (int): How verbose to print, 0 or None = default, 1 = print status every 20 seconds, 2 = also print debugging information
 
     Returns:
         Tuple[Dict[str, Any], Api]: The resolved config and api.
@@ -83,6 +85,7 @@ def resolve_agent_config(  # noqa: C901
         "queues": [],
         "registry": {},
         "builder": {},
+        "verbosity": 0,
     }
     user_set_project = False
     resolved_config: Dict[str, Any] = defaults
@@ -123,6 +126,8 @@ def resolve_agent_config(  # noqa: C901
         resolved_config.update({"max_jobs": int(max_jobs)})
     if queues:
         resolved_config.update({"queues": list(queues)})
+    if verbosity:
+        resolved_config.update({"verbosity": int(verbosity)})
     # queue -> queues
     if resolved_config.get("queue"):
         if isinstance(resolved_config.get("queue"), str):

wandb/sdk/launch/_project_spec.py

@@ -14,6 +14,7 @@ import wandb.docker as docker
 from wandb.apis.internal import Api
 from wandb.errors import CommError
 from wandb.sdk.launch import utils
+from wandb.sdk.launch.utils import get_entrypoint_file
 from wandb.sdk.lib.runid import generate_id
 
 from .errors import LaunchError
@@ -135,7 +136,7 @@ class LaunchProject:
         if override_entrypoint:
             _logger.info("Adding override entry point")
             self.override_entrypoint = EntryPoint(
-                name=_get_entrypoint_file(override_entrypoint),
+                name=get_entrypoint_file(override_entrypoint),
                 command=override_entrypoint,
             )
 
@@ -536,24 +537,6 @@ class LaunchProject:
                 self.git_version = branch_name
 
 
-def _get_entrypoint_file(entrypoint: List[str]) -> Optional[str]:
-    """Get the entrypoint file from the given command.
-
-    Args:
-        entrypoint (List[str]): List of command and arguments.
-
-    Returns:
-        Optional[str]: The entrypoint file if found, otherwise None.
-    """
-    if not entrypoint:
-        return None
-    if entrypoint[0].endswith(".py") or entrypoint[0].endswith(".sh"):
-        return entrypoint[0]
-    if len(entrypoint) < 2:
-        return None
-    return entrypoint[1]
-
-
 class EntryPoint:
     """An entry point into a wandb launch specification."""
 
@@ -570,7 +553,9 @@ class EntryPoint:
 
     def update_entrypoint_path(self, new_path: str) -> None:
         """Updates the entrypoint path to a new path."""
-        if len(self.command) == 2 and self.command[0] in ["python", "bash"]:
+        if len(self.command) == 2 and (
+            self.command[0].startswith("python") or self.command[0] == "bash"
+        ):
             self.command[1] = new_path