mirrors/crosvm
1
0
Fork 0
mirror of https://chromium.googlesource.com/crosvm/crosvm synced 2025-02-10 20:19:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

New merge bot

Browse source 
The new merge bot is written in python to allow us to improve the
merge process.

The new process is as follows:
- The merge bot will only create a new merge when the previous
  ones have been submitted.
- When the bot creates a new merge, it'll split the merge commits
  into smaller changes to make bisection easier.

This allows us to only ever deal with one set of merge commits
instead of piling on new merges day after day.

For dry runs, the new process is:
- The bot will check existing dry runs and abandon them if the CQ
  is done.
- The bot will vote Verified+-1 to indicate if the dry run passed.
- If no more dry runs are active, it'll create a new one.

This ensures we always have a dry run in the CQ and can easily see
in gerrit which ones have passed.

The python script will also allow us to add convenience features
for oncall for retrying bratches, etc.

BUG=b:227347397
TEST=The script can be run locally:
./tools/chromeos/merge_bot update-merges
./tools/chromeos/merge_bot update-dry-runs

Change-Id: I51fc0741f5032868ad07f5e926e570e29cc461fc
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/3583254
Reviewed-by: Junichi Uekawa <uekawa@chromium.org>
Tested-by: kokoro <noreply+kokoro@google.com>
This commit is contained in:
Dennis Kempin 2022-04-12 11:15:20 -07:00
parent a7d91e5f66
commit b193738186
4 changed files with 399 additions and 323 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

217
ci/kokoro/build-merge-into-chromeos.sh
View file

@ -4,219 +4,24 @@
# found in the LICENSE file.
set -e
readonly GERRIT_URL=https://chromium-review.googlesource.com
readonly ORIGIN=https://chromium.googlesource.com/chromiumos/platform/crosvm
readonly RETRIES=3
readonly MIN_COMMIT_COUNT=${MIN_COMMIT_COUNT:-5}
gerrit_api_get() {
# GET request to the gerrit API. Strips XSSI protection line from output.
# See: https://gerrit-review.googlesource.com/Documentation/dev-rest-api.html
local url="${GERRIT_URL}/${1}"
curl --silent "$url" | tail -n +2
}
gerrit_api_post() {
# POST to gerrit API using http cookies from git.
local endpoint=$1
local body=$2
local cookie_file=$(git config http.cookiefile)
if [[ -z "${cookie_file}" ]]; then
echo 1>&2 "Cannot find git http cookie file."
return 1
fi
local url="${GERRIT_URL}/${endpoint}"
curl --silent \
--cookie "${cookie_file}" \
-X POST \
-d "${body}" \
-H "Content-Type: application/json" \
"$url" |
tail -n +2
}
query_change() {
# Query gerrit for a specific change.
# See: https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#get-change
gerrit_api_get "changes/$1/?o=CURRENT_REVISION"
}
query_changes() {
# Query gerrit for a list of changes.
# See: https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#list-changes
local query=$(printf '%s+' "$@")
gerrit_api_get "changes/?q=${query}"
}
query_related_changes() {
# Query related changes from gerrit.
# See: https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#get-related-changes
gerrit_api_get "changes/$1/revisions/current/related"
}
get_previous_merge_id() {
# Query all open merge commits previously made by crosvm-bot. May be null if
# none are open.
query=(
project:chromiumos/platform/crosvm
branch:chromeos
status:open
owner:crosvm-bot@crosvm-packages.iam.gserviceaccount.com
-hashtag:dryrun
)
# Pick the one that was created last.
query_changes "${query[@]}" |
jq --raw-output 'sort_by(.created)[-1].id'
}
get_last_change_in_chain() {
# Use the related changes API to find the last change in the chain of
# commits.
local change_id=$1
# The list of commits is sorted by the git commit order, with the latest
# change first and includes the current change.
local last_change
last_change=$(query_related_changes "$change_id" |
jq --raw-output "[.changes[] | select(.status == \"NEW\")][0].change_id")
# If there are no related changes the list will be empty.
if [ "$last_change" == "null" ]; then
echo "${change_id}"
else
# The related API does not give us the unique ID of changes. Build it manually.
echo "chromiumos%2Fplatform%2Fcrosvm~chromeos~${last_change}"
fi
}
fetch_change() {
# Fetch the provided change and print the commit sha.
local change_id=$1
# Find the git ref we need to fetch.
local change_ref
change_ref=$(query_change "$change_id" |
jq --raw-output -e ".revisions[.current_revision].ref")
git fetch -q origin "${change_ref}"
}
get_dry_run_ids() {
# Query all dry run changes. They are identified by the hashtag:dryrun when
# uploaded.
query=(
project:chromiumos/platform/crosvm
branch:chromeos
status:open
hashtag:dryrun
owner:crosvm-bot@crosvm-packages.iam.gserviceaccount.com
)
query_changes "${query[@]}" |
jq --raw-output '.[].id'
}
abandon_dry_runs() {
# Abandon all pending dry run commits
for change in $(get_dry_run_ids); do
echo "Abandoning ${GERRIT_URL}/q/${change}"
gerrit_api_post "a/changes/${change}/abandon" "{}" >/dev/null
done
}
gerrit_prerequisites() {
# Authenticate to GoB if we don't already have a cookie.
# This should only happen when running in Kokoro, not locally.
# See: go/gob-gce
if [[ -z $(git config http.cookiefile) ]]; then
git clone https://gerrit.googlesource.com/gcompute-tools \
"${KOKORO_ARTIFACTS_DIR}/gcompute-tools"
"${KOKORO_ARTIFACTS_DIR}/gcompute-tools/git-cookie-authdaemon" --no-fork
# Setup correct user info for the service account.
git config user.name "Crosvm Bot"
git config user.email crosvm-bot@crosvm-packages.iam.gserviceaccount.com
fi
# We cannot use the original origin that kokoro used, as we no longer have
# access the GoB host via rpc://.
git remote remove origin
git remote add origin ${ORIGIN}
git fetch -q origin
# Set up gerrit Change-Id hook.
mkdir -p .git/hooks
curl -Lo .git/hooks/commit-msg \
https://gerrit-review.googlesource.com/tools/hooks/commit-msg
chmod +x .git/hooks/commit-msg
}
upload() {
git push origin \
"HEAD:refs/for/chromeos%r=crosvm-uprev@google.com,$1"
}
upload_with_retries() {
# Try uploading to gerrit. Retry due to errors on first upload.
# See: b/209031134
for i in $(seq 1 $RETRIES); do
echo "Push attempt $i"
if upload "$1"; then
return 0
fi
done
return 1
}
main() {
cd "${KOKORO_ARTIFACTS_DIR}/git/crosvm"
gerrit_prerequisites
# Make a copy of the merge script, so we are using the HEAD version to
# create the merge.
cp ./tools/chromeos/create_merge "${KOKORO_ARTIFACTS_DIR}/create_merge"
# Clean possible stray files from previous builds.
git clean -f -d -x
git checkout -f
# Parent commit to use for this merge.
local parent_commit="origin/chromeos"
# Query gerrit to find the latest merge commit and fetch it to be used as
# a parent.
local previous_merge="$(get_previous_merge_id)"
if [ "$previous_merge" != "null" ]; then
# The oncall may have uploaded a custom merge or cherry-pick on top
# of the detected merge. Find the last changed in that chain.
local last_change_in_chain=$(get_last_change_in_chain "${previous_merge}")
echo "Found previous merge: ${GERRIT_URL}/q/${previous_merge}"
echo "Last change in that chain: ${GERRIT_URL}/q/${last_change_in_chain}"
fetch_change "${last_change_in_chain}"
parent_commit="FETCH_HEAD"
# Ensure we have at least python 3.9. Kokoro does not and requires us to use pyenv to install
# The required version.
if ! python -c "import sys \
sys.exit(sys.version_info.major != 3 or sys.version_info.minor < 9)"; then
pyenv install -v 3.9.5
pyenv global 3.9.5
fi
echo "Checking out parent: ${parent_commit}"
git checkout -b chromeos "${parent_commit}"
git branch --set-upstream-to origin/chromeos chromeos
local merge_count=$(git log --oneline --decorate=no --no-color \
"${parent_commit}..origin/main" | wc -l)
if [ "${merge_count}" -ge "$MIN_COMMIT_COUNT" ]; then
"${KOKORO_ARTIFACTS_DIR}/create_merge" "origin/main"
else
echo "Not enough commits to merge."
# Extra packages required by merge_bot
if ! pip show argh; then
pip install argh
fi
upload_with_retries
echo "Abandoning previous dry runs"
abandon_dry_runs
echo "Creating dry run merge"
git checkout -b dryrun --track origin/chromeos
"${KOKORO_ARTIFACTS_DIR}/create_merge" --dry-run-only "origin/main"
upload_with_retries "hashtag=dryrun,l=Commit-Queue+1,l=Bot-Commit+1"
./tools/chromeos/merge_bot -v update-merges --is-bot
./tools/chromeos/merge_bot -v update-dry-runs --is-bot
}
main

112
tools/chromeos/create_merge
View file

@ -1,112 +0,0 @@
#!/bin/bash
# Copyright 2021 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
#
# Script to create a commit to merge cros/main into cros/chromeos with a useful
# commit message.
#
# Basic usage to upload a merge to gerrit:
#
# $ repo start uprev .
# $ ./tools/chromeos/create_merge
# $ git push cros HEAD:refs/for/chromeos
#
# To merge with a specific commit, use: ./tools/chromeos/create_merge $SHA
set -e
if [ "$1" == "--dry-run-only" ]; then
DRY_RUN_ONLY=true
shift
fi
LOCAL_BRANCH=$(git branch --show-current)
REMOTE_NAME=$(git config "branch.${LOCAL_BRANCH}.remote")
URL=$(git remote get-url "${REMOTE_NAME}")
DEFAULT_TARGET="${REMOTE_NAME}/main"
MERGE_TARGET="${1:-${DEFAULT_TARGET}}"
commit_list() {
git log --oneline --decorate=no --no-color "HEAD..${MERGE_TARGET}"
}
prerequisites() {
if [[ -e "${LOCAL_BRANCH}" ]] ||
[[ -e "${REMOTE_NAME}" ]] ||
[[ -e "${URL}" ]]; then
echo "This script requires the local repository to be on" \
"a tracking branch."
exit 1
fi
if [[ -n $(git status -s) ]]; then
echo "Working directory is not clean:"
git status -s
exit 1
fi
if [[ -z "$(commit_list)" ]]; then
echo "Nothing to merge."
exit 0
fi
}
cq_depends() {
git log --no-color "HEAD..${MERGE_TARGET}" --pretty=email |
grep ^Cq-Depend: |
sort -u
}
bug_references() {
git log --no-color "HEAD..${MERGE_TARGET}" --pretty=email |
grep ^BUG= |
grep -vi ^BUG=none |
sort -u
}
merge_message() {
local old=$(git rev-parse HEAD)
local new=$(git rev-parse "${MERGE_TARGET}")
local count=$(commit_list | wc -l)
local notes="$(date +%F)"
if [[ -n "$(cq_depends)" ]]; then
notes="${notes}, cq-depend"
fi
if [ "${DRY_RUN_ONLY}" = true ]; then
echo "Merge dry run (${notes})"
else
echo "Merge ${count} commits from ${MERGE_TARGET} (${notes})"
fi
echo ""
commit_list
echo ""
echo "${URL}/+log/${old}..${new}"
echo ""
if [ "${DRY_RUN_ONLY}" != true ]; then
bug_references
fi
echo "TEST=CQ"
}
merge_trailers() {
cq_depends
if [ "${DRY_RUN_ONLY}" = true ]; then
echo "Commit: false"
fi
}
main() {
prerequisites
# Note: trailers need to be added in a separate -m argument. Otherwise trailing whitespace may
# be trimmed which can confuse the gerrit preupload hook when it's trying to add the Commit-Id
# trailer.
git merge -X theirs --no-ff "${MERGE_TARGET}" -m "$(merge_message)" -m "$(merge_trailers)"
git --no-pager log -n 1
}
main

365
tools/chromeos/merge_bot Executable file
View file

@ -0,0 +1,365 @@
#!/usr/bin/env python3
# Copyright 2022 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
# This script is used by the CI system to regularly update the merge and dry run changes.
#
# It can be run locally as well, however some permissions are only given to the bot's service
# account (and are enabled with --is-bot).
#
# See `./tools/chromeos/merge_bot -h` for details.
import functools
import os
import sys
from datetime import date
from typing import Any, cast
sys.path.append(os.path.dirname(sys.path[0]))
import json
import re
from impl.common import CROSVM_ROOT, batched, cmd, quoted, run_commands, very_verbose
git = cmd("git")
git_log = git("log --decorate=no --color=never")
curl = cmd("curl --silent --fail-with-body")
chmod = cmd("chmod")
GERRIT_URL = "https://chromium-review.googlesource.com"
REPO_URL = "https://chromium.googlesource.com/chromiumos/platform/crosvm"
# Gerrit tags used to identify bot changes.
MERGE_TAG = "crosvm-merge"
DRY_RUN_TAG = "crosvm-merge-dry-run"
# This is the email of the account that posts CQ messages.
LUCI_EMAIL = "chromeos-scoped@luci-project-accounts.iam.gserviceaccount.com"
def is_gce_instance():
"""
Returns true if we are running in a GCE instance.
See http://cloud.google.com/compute/docs/instances/detect-compute-engine
"""
try:
metadata = curl("http://metadata.google.internal").stdout()
return "Metadata-Flavor: Google" in metadata
except:
return False
def strip_xssi(response: str):
# See https://gerrit-review.googlesource.com/Documentation/rest-api.html#output
assert response.startswith(")]}'\n")
return response[5:]
def gerrit_api_get(path: str):
response = curl(f"{GERRIT_URL}/{path}").stdout()
return json.loads(strip_xssi(response))
def gerrit_api_post(path: str, body: Any):
cookiefile = git("config http.cookiefile").stdout()
response = curl(
"--cookie",
cookiefile,
"-X POST",
"-H",
quoted("Content-Type: application/json"),
"-d",
quoted(json.dumps(body)),
f"{GERRIT_URL}/{path}",
).stdout()
if very_verbose():
print("Response:", response)
return json.loads(strip_xssi(response))
class Change(object):
"""
Class to interact with the gerrit /changes/ API.
For information on the data format returned by the API, see:
https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#change-info
"""
id: str
_data: Any
def __init__(self, data: Any):
self._data = data
self.id = data["id"]
@functools.cached_property
def _details(self) -> Any:
return gerrit_api_get(f"changes/{self.id}/detail")
@functools.cached_property
def _messages(self) -> list[Any]:
return gerrit_api_get(f"changes/{self.id}/messages")
def get_votes(self, label_name: str) -> list[int]:
"Returns the list of votes on `label_name`"
label_info = self._details.get("labels", {}).get(label_name)
votes = label_info.get("all", [])
return [cast(int, v.get("value")) for v in votes]
def get_messages_by(self, email: str) -> list[str]:
"Returns all messages posted by the user with the specified `email`."
return [m["message"] for m in self._messages if m["author"].get("email") == email]
def review(self, message: str, labels: dict[str, int]):
"Post review `message` and set the specified review `labels`"
print("Posting on", self, ":", message, labels)
gerrit_api_post(
f"changes/{self.id}/revisions/current/review",
{"message": message, "labels": labels},
)
def abandon(self, message: str):
print("Abandoning", self, ":", message)
gerrit_api_post(f"changes/{self.id}/abandon", {"message": message})
@staticmethod
def query(*queries: str):
"Returns a list of gerrit changes matching the provided list of queries."
return [Change(c) for c in gerrit_api_get(f"changes/?q={'+'.join(queries)}")]
def __str__(self):
return f"http://crrev.com/c/{self._data['_number']}"
def pretty_info(self):
return f"{self} - {self._data['subject']}"
def list_active_merges():
return Change.query(
"project:chromiumos/platform/crosvm",
"branch:chromeos",
"status:open",
f"hashtag:{MERGE_TAG}",
)
def list_active_dry_runs():
return Change.query(
"project:chromiumos/platform/crosvm",
"branch:chromeos",
"status:open",
f"hashtag:{DRY_RUN_TAG}",
)
def bug_notes(commit_range: str):
"Returns a string with all BUG=... lines of the specified commit range."
return "\n".join(
set(
line
for line in git_log(commit_range, "--pretty=%b").lines()
if re.match(r"^BUG=", line, re.I) and not re.match(r"^BUG=None", line, re.I)
)
)
def setup_tracking_branch(branch_name: str, tracking: str):
"Create and checkout `branch_name` tracking `tracking`. Overwrites existing branch."
git("checkout", tracking).fg(quiet=True)
git("branch -D", branch_name).fg(quiet=True, check=False)
git("checkout -b", branch_name, "--track", tracking).fg()
def gerrit_prerequisites():
"Make sure we can upload to gerrit."
# Make sure we have http cookies to access gerrit
cookie_file = git("config http.cookiefile").stdout()
if not cookie_file:
if is_gce_instance():
# Grab http cookies for accessing GOB. See go/gob-gce
gcompute_path = f"{os.environ['KOKORO_ARTIFACTS_DIR']}/gcompute-tools"
git("clone", "https://gerrit.googlesource.com/gcompute-tools", gcompute_path)
cmd(f"{gcompute_path}/git-cookie-authdaemon", "--no-fork").fg()
# Setup correct user info for the GCE service account.
git("config user.name", quoted("Crosvm Bot"))
git("config user.email", quoted("crosvm-bot@crosvm-packages.iam.gserviceaccount.com"))
else:
print("Cannot access gerrit without http cookies.")
print(f"Install one via: {GERRIT_URL}/new-password")
sys.exit(1)
# Re-add remote to ensure it is correct. We cannot use the 'origin' provided by Kokoro.
git("remote remove origin").fg()
git("remote add origin", REPO_URL).fg()
git("fetch -q origin").fg()
# Install gerrit Change-Id hook
hook_path = CROSVM_ROOT / ".git/hooks/commit-msg"
if not hook_path.exists():
hook_path.parent.mkdir(exist_ok=True)
curl(f"{GERRIT_URL}/tools/hooks/commit-msg").write_to(hook_path)
chmod("+x", hook_path).fg()
def upload_to_gerrit(target_branch: str, *extra_params: str):
extra_params = ("r=crosvm-uprev@google.com", *extra_params)
for i in range(3):
try:
print(f"Uploading to gerrit (Attempt {i})")
git(f"push origin HEAD:refs/for/{target_branch}%{','.join(extra_params)}").fg()
break
except:
continue
####################################################################################################
# The functions below are callable via the command line
def create_merge_commits(
revision: str, max_size: int = 0, create_dry_run: bool = False, is_bot: bool = False
):
"Merges `revision` into HEAD, creating merge commits including at most `max-size` commits."
os.chdir(CROSVM_ROOT)
from_rev = git("rev-parse", revision).stdout()
# Find list of commits to merge, then batch them into smaller merges.
commits = git_log(f"HEAD..{from_rev}", "--pretty=%H").lines()
if not commits:
print("Nothing to merge.")
return 0
# Create a merge commit for each batch
batches = list(batched(commits, max_size)) if max_size > 0 else [commits]
for i, batch in enumerate(reversed(batches)):
target = batch[0]
commit_range = f"{batch[-1]}..{batch[0]}"
# Put together a message containing info about what's in the merge.
batch_str = f"{i + 1}/{len(batches)}" if len(batches) > 1 else ""
title = "Merge with main" if not create_dry_run else f"Merge dry run"
message = "\n\n".join(
[
f"{title} {date.today().isoformat()} {batch_str}",
git_log(commit_range, "--oneline").stdout(),
f"{REPO_URL}/+log/{commit_range}",
bug_notes(commit_range),
]
)
# git 'trailers' go into a separate paragraph to make sure they are properly separated.
trailers = "Commit: False" if create_dry_run else ""
# Perfom merge
git("merge --no-ff", target, "-m", quoted(message), "-m", quoted(trailers)).fg()
return len(batches)
def status():
"Shows the current status of pending merge and dry run changes in gerrit."
print("Active dry runs:")
for dry_run in list_active_dry_runs():
print(dry_run.pretty_info())
print()
print("Active merges:")
for merge in list_active_merges():
print(merge.pretty_info())
def update_merges(
revision: str = "origin/main",
target_branch: str = "chromeos",
max_size: int = 15,
is_bot: bool = False,
):
"""Uploads a new set of merge commits if the previous batch has been submitted."""
gerrit_prerequisites()
active_merges = list_active_merges()
if active_merges:
print("Nothing to do. Previous merges are still pending:")
for merge in active_merges:
print(merge.pretty_info())
return
else:
print(f"Creating merge of {revision} into origin/{target_branch}")
setup_tracking_branch("merge-bot-branch", f"origin/{target_branch}")
if create_merge_commits(revision, max_size, create_dry_run=False) > 0:
upload_to_gerrit(
target_branch,
f"hashtag={MERGE_TAG}",
"l=Commit-Queue+1",
*(["l=Bot-Commit+1"] if is_bot else []),
)
def update_dry_runs(
revision: str = "origin/main",
target_branch: str = "chromeos",
max_size: int = 15,
is_bot: bool = False,
):
"""
Maintains dry run changes in gerrit, usually run by the crosvm bot, but can be called by
developers as well.
"""
gerrit_prerequisites()
# Close active dry runs if they are done.
print("Checking active dry runs")
for dry_run in list_active_dry_runs():
if max(dry_run.get_votes("Commit-Queue")) > 0:
print(dry_run, "CQ is still running.")
continue
# Check for luci results and add V+-1 votes to make it easier to identify failed dry runs.
luci_messages = dry_run.get_messages_by(LUCI_EMAIL)
if not luci_messages:
print(dry_run, "No luci messages yet.")
continue
last_luci_message = luci_messages[-1]
if "This CL passed the CQ dry run" in last_luci_message:
dry_run.review(
"I think this dry run was SUCCESSFUL.",
{
"Verified": 1,
"Bot-Commit": 0,
},
)
elif "Failed builds" in last_luci_message:
dry_run.review(
"I think this dry run FAILED.",
{
"Verified": -1,
"Bot-Commit": 0,
},
)
dry_run.abandon("I am creating a new dry run.")
active_dry_runs = list_active_dry_runs()
if active_dry_runs:
print("There are active dry runs, not creating a new one.")
print("Active dry runs:")
for dry_run in active_dry_runs:
print(dry_run.pretty_info())
return
print(f"Creating dry run merge of {revision} into origin/{target_branch}")
setup_tracking_branch("merge-bot-branch", f"origin/{target_branch}")
if create_merge_commits(revision, max_size, create_dry_run=True) > 0:
upload_to_gerrit(
target_branch,
f"hashtag={DRY_RUN_TAG}",
"l=Commit-Queue+1",
*(["l=Bot-Commit+1"] if is_bot else []),
)
run_commands(create_merge_commits, status, update_merges, update_dry_runs)

28
tools/impl/common.py
View file

@ -17,6 +17,7 @@ from __future__ import annotations
import argparse
import contextlib
import csv
from math import ceil
import os
import re
import shutil
@ -173,7 +174,7 @@ class Command(object):
)
if result.returncode != 0:
if quiet and result.stdout:
if quiet and check and result.stdout:
print(result.stdout)
if check:
raise subprocess.CalledProcessError(result.returncode, str(self), result.stdout)
@ -185,6 +186,12 @@ class Command(object):
"""
return self.run(stderr=None).stdout.strip()
def lines(self):
"""
Runs a program and returns stdout line by line. Stderr is still directed to the user.
"""
return self.stdout().splitlines()
def write_to(self, filename: Path):
"""
Writes all program output (stdout and stderr) to the provided file.
@ -413,7 +420,7 @@ class QuotedString(object):
T = TypeVar("T")
def batched(source: Iterable[T], batch_size: int) -> Iterable[list[T]]:
def batched(source: Iterable[T], max_batch_size: int) -> Iterable[list[T]]:
"""
Returns an iterator over batches of elements from source_list.
@ -421,6 +428,9 @@ def batched(source: Iterable[T], batch_size: int) -> Iterable[list[T]]:
[[1, 2], [3, 4], [5]]
"""
source_list = list(source)
# Calculate batch size that spreads elements evenly across all batches
batch_count = ceil(len(source_list) / max_batch_size)
batch_size = ceil(len(source_list) / batch_count)
for index in range(0, len(source_list), batch_size):
yield source_list[index : min(index + batch_size, len(source_list))]
@ -433,16 +443,24 @@ parallel = ParallelCommands
def run_main(main_fn: Callable[..., Any]):
run_commands(default_fn=main_fn)
def run_commands(*functions: Callable[..., Any], default_fn: Optional[Callable[..., Any]] = None):
"""
Runs the main function using argh to translate command line arguments into function arguments.
Allow the user to call the provided functions with command line arguments translated to
function arguments via argh: https://pythonhosted.org/argh
"""
try:
# Add global verbose arguments
parser = argparse.ArgumentParser()
__add_verbose_args(parser)
# Register main method as argh command
argh.set_default_command(parser, main_fn) # type: ignore
# Add provided commands to parser. Do not use sub-commands if we just got one function.
if functions:
argh.add_commands(parser, functions) # type: ignore
if default_fn:
argh.set_default_command(parser, default_fn) # type: ignore
# Call main method
argh.dispatch(parser) # type: ignore