Skip to content

Commit

Permalink
Add docstrings to new HUB functions (#7576)
Browse files Browse the repository at this point in the history
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
  • Loading branch information
glenn-jocher committed Jan 14, 2024
1 parent d762496 commit 2f11ab5
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions ultralytics/hub/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ def __init__(self, identifier):
self.model = self.client.model() # load empty model

def load_model(self, model_id):
# Initialize model
"""Loads an existing model from Ultralytics HUB using the provided model identifier."""
self.model = self.client.model(model_id)
if not self.model.data: # then model model does not exist
raise ValueError(emojis(f"❌ The specified HUB model does not exist")) # TODO: improve error handling
Expand All @@ -82,7 +82,7 @@ def load_model(self, model_id):
LOGGER.info(f"{PREFIX}View model at {self.model_url} 🚀")

def create_model(self, model_args):
# Initialize model
"""Initializes a HUB training session with the specified model identifier."""
payload = {
"config": {
"batchSize": model_args.get("batch", -1),
Expand Down Expand Up @@ -168,6 +168,7 @@ def _parse_identifier(self, identifier):
return api_key, model_id, filename

def _set_train_args(self, **kwargs):
"""Initializes training arguments and creates a model entry on the Ultralytics HUB."""
if self.model.is_trained():
# Model is already trained
raise ValueError(emojis(f"Model is already trained and uploaded to {self.model_url} 🚀"))
Expand All @@ -179,6 +180,7 @@ def _set_train_args(self, **kwargs):
else:
# Model has no saved weights
def get_train_args(config):
"""Parses an identifier to extract API key, model ID, and filename if applicable."""
return {
"batch": config["batchSize"],
"epochs": config["epochs"],
Expand Down Expand Up @@ -213,6 +215,7 @@ def request_queue(
**kwargs,
):
def retry_request():
"""Attempts to call `request_func` with retries, timeout, and optional threading."""
t0 = time.time() # Record the start time for the timeout
for i in range(retry + 1):
if (time.time() - t0) > timeout:
Expand Down Expand Up @@ -254,7 +257,7 @@ def retry_request():
return retry_request()

def _should_retry(self, status_code):
# Status codes that trigger retries
"""Determines if a request should be retried based on the HTTP status code."""
retry_codes = {
HTTPStatus.REQUEST_TIMEOUT,
HTTPStatus.BAD_GATEWAY,
Expand Down

0 comments on commit 2f11ab5

Please sign in to comment.