diff --git a/benchmarking/frameworks/framework_base.py b/benchmarking/frameworks/framework_base.py index 7c144224..49aa53dc 100644 --- a/benchmarking/frameworks/framework_base.py +++ b/benchmarking/frameworks/framework_base.py @@ -20,6 +20,7 @@ import random import re import shutil +from copy import deepcopy from bridge.file_storage.upload_files.file_uploader import FileUploader from data_converters.data_converters import getConverters @@ -547,16 +548,30 @@ def _runCommands( main_command, ) profiling_enabled = False + profiling_args = {} if "profiler" in test: - profiling_enabled = test.get("profiler", {}).get("enabled", False) + profiling_enabled = test["profiler"].get("enabled", False) if profiling_enabled: - platform_args["profiler_args"] = test.get("profiler", {}) + # test[] is potentially raw user input so we need to ensure + # ensure all fields are populated so we don't have to check elsewhere + profiling_args = deepcopy(test["profiler"]) + default_profiler = ( + "perfetto" + if "cpu" not in profiling_args.get("types", ["cpu"]) + else "simpleperf" + ) + profiler = profiling_args.setdefault("profiler", default_profiler) + default_type = "memory" if profiler == "perfetto" else "cpu" + profiling_args.setdefault("types", [default_type]) + profiling_args.setdefault("options", {}) platform_args["model_name"] = getModelName(model) for idx, cmd in enumerate(cmds): # note that we only enable profiling for the last command # of the main commands. - platform_args["enable_profiling"] = ( - profiling_enabled and main_command and idx == len(cmds) - 1 + platform_args["profiling_args"] = ( + profiling_args + if (profiling_enabled and main_command and idx == len(cmds) - 1) + else {"enabled": False} ) one_output = self.runOnPlatform( total_num, cmd, platform, platform_args, converter diff --git a/benchmarking/platforms/android/android_platform.py b/benchmarking/platforms/android/android_platform.py index cc25c030..288b5c95 100644 --- a/benchmarking/platforms/android/android_platform.py +++ b/benchmarking/platforms/android/android_platform.py @@ -17,6 +17,7 @@ from degrade.degrade_base import DegradeBase, getDegrade from platforms.platform_base import PlatformBase +from profilers.perfetto.perfetto import Perfetto from profilers.profilers import getProfilerByUsage from six import string_types from utils.custom_logger import getLogger @@ -71,7 +72,7 @@ def _setLogCatSize(self): # We know this command may fail. Avoid propogating this # failure to the upstream success = getRunStatus() - ret = self.util.run(["logcat", "-G", str(size) + "K"], timeout=2, retry=1) + ret = self.util.logcat("-G", str(size) + "K") setRunStatus(success, overwrite=True) if len(ret) > 0 and ret[0].find("failed to") >= 0: repeat = True @@ -213,7 +214,6 @@ def runBinaryBenchmark(self, cmd, *args, **kwargs): "log_to_screen_only" in kwargs and kwargs["log_to_screen_only"] ) platform_args = {} - meta = {} if "platform_args" in kwargs: platform_args = kwargs["platform_args"] if "taskset" in platform_args: @@ -236,38 +236,96 @@ def runBinaryBenchmark(self, cmd, *args, **kwargs): ) platform_args["non_blocking"] = True del platform_args["power"] - if platform_args.get("enable_profiling", False): - # attempt to run with profiling, else fallback to standard run - try: - simpleperf = getProfilerByUsage( - "android", - None, - platform=self, - model_name=platform_args.get("model_name", None), - cmd=cmd, - ) - if simpleperf: - f = simpleperf.start() - output, meta = f.result() - if not output or not meta: - raise RuntimeError( - "No data returned from Simpleperf profiler." - ) - log_logcat = [] - if not log_to_screen_only: - log_logcat = self.util.logcat("-d") - return output + log_logcat, meta - # if this has not succeeded for some reason reset run status and run without profiling. - except Exception: - getLogger().critical( - f"An error has occurred when running Simpleperf profiler on device {self.platform} {self.platform_hash}.", - exc_info=True, + enable_profiling = platform_args.get("profiling_args", {}).get( + "enabled", False + ) + if enable_profiling: + profiler = platform_args["profiling_args"]["profiler"] + profiling_types = platform_args["profiling_args"]["types"] + if profiler == "simpleperf": + assert profiling_types == [ + "cpu" + ], "Only cpu profiling is supported for SimplePerf" + try: + # attempt to run with cpu profiling, else fallback to standard run + return self._runBenchmarkWithSimpleperf( + cmd, log_to_screen_only, **platform_args + ) + except Exception: + # if this has not succeeded for some reason reset run status and run without profiling. + getLogger().critical( + f"An error has occurred when running Simpleperf profiler on device {self.platform} {self.platform_hash}.", + exc_info=True, + ) + elif profiler == "perfetto": + assert ( + "cpu" not in profiling_types + ), "cpu profiling is not yet implemented for Perfetto" + try: + # attempt Perfetto profiling + return self._runBenchmarkWithPerfetto( + cmd, log_to_screen_only, **platform_args + ) + except Exception: + # if this has not succeeded for some reason reset run status and run without profiling. + getLogger().critical( + f"An error has occurred when running Perfetto profiler on device {self.platform} {self.platform_hash}.", + exc_info=True, + ) + else: + getLogger().error( + f"Ignoring unsupported profiler setting: {profiler}: {profiling_types}.", ) + + # Run without profiling + return self._runBinaryBenchmark(cmd, log_to_screen_only, **platform_args) + + def _runBinaryBenchmark(self, cmd, log_to_screen_only: bool, **platform_args): log_screen = self.util.shell(cmd, **platform_args) log_logcat = [] if not log_to_screen_only: log_logcat = self.util.logcat("-d") - return log_screen + log_logcat, meta + return log_screen + log_logcat, {} + + def _runBenchmarkWithSimpleperf( + self, cmd, log_to_screen_only: bool, **platform_args + ): + simpleperf = getProfilerByUsage( + "android", + None, + platform=self, + model_name=platform_args.get("model_name", None), + cmd=cmd, + ) + if simpleperf: + f = simpleperf.start() + output, meta = f.result() + if not output or not meta: + raise RuntimeError("No data returned from Simpleperf profiler.") + log_logcat = [] + if not log_to_screen_only: + log_logcat = self.util.logcat("-d") + return output + log_logcat, meta + + def _runBenchmarkWithPerfetto(self, cmd, log_to_screen_only: bool, **platform_args): + # attempt Perfetto profiling + if not self.util.isRootedDevice(silent=True): + raise RuntimeError( + "Attempted to perform Perfetto profiling on unrooted device {self.util.device}." + ) + + with Perfetto( + platform=self, + types=platform_args["profiling_args"]["types"], + options=platform_args["profiling_args"]["options"], + ) as perfetto: + getLogger().info("Invoked with Perfetto.") + log_screen = self.util.shell(cmd, **platform_args) + log_logcat = [] + if not log_to_screen_only: + log_logcat = self.util.logcat("-d") + meta = perfetto.getResults() + return log_screen + log_logcat, meta def collectMetaData(self, info): meta = super(AndroidPlatform, self).collectMetaData(info) diff --git a/benchmarking/profilers/perfetto/__init__.py b/benchmarking/profilers/perfetto/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/benchmarking/profilers/perfetto/perfetto.py b/benchmarking/profilers/perfetto/perfetto.py new file mode 100644 index 00000000..ced00689 --- /dev/null +++ b/benchmarking/profilers/perfetto/perfetto.py @@ -0,0 +1,410 @@ +#!/usr/bin/env python + +############################################################################## +# Copyright 2021-present, Facebook, Inc. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +############################################################################## + +import logging +import os +import shutil +import tempfile +import time +from pathlib import Path +from tempfile import NamedTemporaryFile +from typing import Optional + +# from platforms.android.android_platform import AndroidPlatform +from profilers.perfetto.perfetto_config import ( + CONFIG_TEMPLATE, + POWER_CONFIG, + HEAPPROFD_CONFIG, + ANDROID_LOG_CONFIG, + LINUX_FTRACE_CONFIG, +) +from profilers.profiler_base import ProfilerBase +from profilers.utilities import generate_perf_filename, upload_profiling_reports +from utils.custom_logger import getLogger + +PROCESS_KEY = "perfetto" + +""" +Perfetto is a native memory and battery profiling tool for Android OS 10 or better. +It can be used to profile both Android applications and native +processes running on Android. It can profile both Java and C++ code on Android. + +Perfetto can be used to profile Android benchmarks as both applications and +binaries. The resulting perf data is used to generate an html report +including a flamegraph (TODO). Both perf data and the report are uploaded to manifold +and the urls are returned as a meta dict which can be updated in the benchmark's meta data. +""" + +logger = logging.getLogger(__name__) + + +class Perfetto(ProfilerBase): + + CONFIG_FILE = "perfetto.conf" + DEVICE_DIRECTORY = "/data/local/tmp/perf" + TRACING_PROPERTY = "persist.traced.enable" + DEFAULT_TIMEOUT = 5 + BUFFER_SIZE_KB_DEFAULT = 256 * 1024 # 256 megabytes + BUFFER_SIZE2_KB_DEFAULT = 2 * 1024 # 2 megabytes + SHMEM_SIZE_BYTES_DEFAULT = 8388608 + SAMPLING_INTERVAL_BYTES_DEFAULT = 4096 + BATTERY_POLL_MS_DEFAULT = 1000 + MAX_FILE_SIZE_BYTES_DEFAULT = 100000000 + + def __init__( + self, + platform, + *, + types=None, + options=None, + model_name="benchmark", + ): + self.platform = platform + self.types = types or ["memory"] + self.options = options or {} + self.android_version: int = int(platform.rel_version.split(".")[0]) + self.adb = platform.util + self.valid = True + self.perfetto_pid = None + self.all_heaps = ( + f"all_heaps: {self.options.get('all_heaps', 'false')}" + if self.android_version >= 12 + else "" + ) + self.basename = generate_perf_filename(model_name, self.adb.device) + self.trace_file_name = f"{self.basename}.perfetto-trace" + self.trace_file_device = f"{self.DEVICE_DIRECTORY}/{self.trace_file_name}" + self.config_file = f"{self.basename}.{self.CONFIG_FILE}" + self.config_file_device = f"{self.DEVICE_DIRECTORY}/{self.config_file}" + self.data_file = f"{self.basename}.data.json" + self.report_file = f"{self.basename}.txt" # f"{self.basename}.html" + self.user_home = str(Path.home()) + self.host_binary_location = f"{self.user_home}/android" + self.host_output_dir = "" + self.meta = {} + self.is_rooted_device = self.adb.isRootedDevice() + self.user_was_root = self.adb.user_is_root() if self.is_rooted_device else False + self.original_SELinux_policy = ( + self.adb.shell( + ["getenforce"], + default=[""], + timeout=self.DEFAULT_TIMEOUT, + )[0] + .strip() + .lower() + ) + self.perfetto_cmd = [ + "perfetto", + "-d", + "--txt", + "-c", + self.config_file_device, + "-o", + self.trace_file_device, + ] + super(Perfetto, self).__init__(None) + + def __enter__(self): + self._start() + + return self + + def __exit__(self, type, value, traceback): + if self.meta == {}: + self.meta = self._finish() + + def _start(self): + """Begin Perfetto profiling on platform.""" + try: + if self.android_version < 10: + getLogger().error( + f"Attempt to run Perfetto on {self.platform.type} {self.platform.rel_version} device {self.adb.device} ignored." + ) + self.valid = False + return None + + if not self.is_rooted_device: + getLogger().error( + f"Attempt to run Perfetto on unrooted device {self.adb.device} ignored." + ) + self.valid = False + return None + + getLogger().info(f"Collect Perfetto data on device {self.adb.device}") + self._enablePerfetto() + + # Generate and upload custom config file + getLogger().info(f"Perfetto profile type(s) = {','.join(self.types)}.") + self._setup_perfetto_config() + """ + # Ensure no old instances of perfetto are running on the device + self.adb.shell( + ["killall", "perfetto"], + timeout=DEFAULT_TIMEOUT, + ) + """ + + # call Perfetto + output = self._perfetto() + if output != 1 and output[0] != "1": + self.perfetto_pid = output[0] + return output + except Exception: + self.valid = False + getLogger().exception("Perfetto profiling could not be started.") + return None + + def getResults(self): + if self.valid: + self.meta = self._finish() + + return self.meta + + def _finish(self): + no_report_str = "Perfetto profiling reporting could not be completed." + if not self.valid: + self._restoreState() + return {} + + meta = {} + self.host_output_dir = tempfile.mkdtemp() + try: + # if we ran perfetto, signal it to stop profiling + if self._signalPerfetto(): + getLogger().info( + f"Looking for Perfetto data on device {self.adb.device}" + ) + self._copyPerfDataToHost() + self._generateReport() + meta = self._uploadResults() + else: + getLogger().error( + no_report_str, + ) + except Exception as e: + getLogger().exception( + no_report_str + f" {e}", + exc_info=True, + ) + + # TODO: remove reboot + sleep once this is done in device manager + self.adb.reboot() + time.sleep(10) + meta = {} + finally: + self._restoreState() + shutil.rmtree(self.host_output_dir) + self.valid = False # prevent additional calls + + return meta + + def _uploadResults(self): + meta = upload_profiling_reports( + { + "perfetto_config": os.path.join(self.host_output_dir, self.config_file), + "perfetto_data": os.path.join( + self.host_output_dir, self.trace_file_name + ), + # TODO: generate flamegraph here + "perfetto_report": os.path.join(self.host_output_dir, self.config_file), + } + ) + getLogger().info( + f"Perfetto profiling data uploaded.\nPerfetto Config:\t{meta['perfetto_config']}\nPerfetto Data: \t{meta['perfetto_data']}\nPerfetto Report:\t{meta['perfetto_report']}" + ) + + return meta + + def _restoreState(self): + if self.original_SELinux_policy == "enforcing": + self.adb.shell( + ["setenforce", "1"], + timeout=self.DEFAULT_TIMEOUT, + retry=1, + ) + if (not self.user_was_root) and self.adb.user_is_root(): + self.adb.unroot() # unroot only if it was not rooted to start + + def _signalPerfetto(self) -> bool: + # signal perfetto to stop profiling and await results + getLogger().info("Stopping Perfetto profiling.") + result = None + if self.perfetto_pid is not None: + sigint_cmd = [ + "kill", + "-SIGINT", + self.perfetto_pid, + "&&", + "wait", + self.perfetto_pid, + ] + sigterm_cmd = ["kill", "-SIGTERM", self.perfetto_pid] + else: + sigint_cmd = ["pkill", "-SIGINT", "perfetto"] + sigterm_cmd = ["pkill", "-SIGTERM", "perfetto"] + + cmd = sigint_cmd + try: + # Wait for Perfetto to finish gracefully + getLogger().info("Running '" + " ".join(cmd) + "'.") + result = self.adb.shell( + sigint_cmd, + timeout=30, + retry=1, + silent=True, + ) + if self.perfetto_pid is None: + time.sleep(6.0) + return True + except Exception as e: + getLogger().exception( + f"Perfetto did not respond to SIGINT. Terminating. {e}." + ) + cmd = sigterm_cmd + result = self.adb.shell( + cmd, + timeout=10, + ) + return False + finally: + getLogger().info(f"Running '{' '.join(cmd)}' returned {result}.") + + def _enablePerfetto(self): + if not self.user_was_root: + self.adb.root() + + # Set SELinux to permissive mode if not already + if self.original_SELinux_policy == "enforcing": + self.adb.shell( + ["setenforce", "0"], + timeout=self.DEFAULT_TIMEOUT, + retry=1, + ) + + # Enable Perfetto if not enabled yet. + getprop_tracing_enabled = self.adb.getprop( + self.TRACING_PROPERTY, + default=["0"], + timeout=self.DEFAULT_TIMEOUT, + ) + perfetto_enabled: str = ( + getprop_tracing_enabled if getprop_tracing_enabled else "0" + ) + if not perfetto_enabled.startswith("1"): + self.adb.setprop( + self.TRACING_PROPERTY, + "1", + timeout=self.DEFAULT_TIMEOUT, + ) + + def _setup_perfetto_config( + self, + *, + app_name: str = "program", + config_file_host: Optional[str] = None, + android_logcat: bool = False, + ): + with NamedTemporaryFile() as f: + if config_file_host is None: + # Write custom perfetto config + config_file_host = f.name + heapprofd_config = "" + power_config = "" + linux_process_stats_config = "" + linux_ftrace_config = "" + android_log_config = "" + track_event_config = "" + buffer_size_kb = self.options.get( + "buffer_size_kb", self.BUFFER_SIZE_KB_DEFAULT + ) + buffer_size2_kb = self.options.get( + "buffer_size2_kb", self.BUFFER_SIZE2_KB_DEFAULT + ) + max_file_size_bytes = self.options.get( + "max_file_size_bytes", self.MAX_FILE_SIZE_BYTES_DEFAULT + ) + if "memory" in self.types: + shmem_size_bytes = self.options.get( + "shmem_size_bytes", self.SHMEM_SIZE_BYTES_DEFAULT + ) + sampling_interval_bytes = self.options.get( + "sampling_interval_bytes", self.SAMPLING_INTERVAL_BYTES_DEFAULT + ) + heapprofd_config = HEAPPROFD_CONFIG.format( + all_heaps=self.all_heaps, + shmem_size_bytes=shmem_size_bytes, + sampling_interval_bytes=sampling_interval_bytes, + app_name=app_name, + ) + if "battery" in self.types: + battery_poll_ms = self.options.get( + "battery_poll_ms", self.BATTERY_POLL_MS_DEFAULT + ) + power_config = POWER_CONFIG.format( + battery_poll_ms=battery_poll_ms, + ) + linux_ftrace_config = LINUX_FTRACE_CONFIG.format( + app_name=app_name, + ) + + if "cpu" in self.types: + getLogger().error( + "Error: CPU profiling with perfetto is Not Yet Implemented.", + ) + + if android_logcat: + android_log_config = ANDROID_LOG_CONFIG + + # Generate config file + config_str = CONFIG_TEMPLATE.format( + max_file_size_bytes=max_file_size_bytes, + buffer_size_kb=buffer_size_kb, + buffer_size2_kb=buffer_size2_kb, + android_log_config=android_log_config, + power_config=power_config, + heapprofd_config=heapprofd_config, + linux_process_stats_config=linux_process_stats_config, + linux_ftrace_config=linux_ftrace_config, + track_event_config=track_event_config, + ) + f.write(config_str.encode("utf-8")) + f.flush() + + # Push perfetto config to device + getLogger().info( + f"Host config file = {config_file_host},\nDevice config file = {self.config_file_device}." + ) + self.adb.push(config_file_host, self.config_file_device) + + # Setup permissions for it, to avoid perfetto call failure + self.adb.shell(["chmod", "777", self.config_file_device]) + + def _perfetto(self): + """Run perfetto on platform with benchmark process id.""" + getLogger().info(f"Calling Perfetto: {self.perfetto_cmd}") + output = self.platform.util.shell(self.perfetto_cmd) + getLogger().info(f"Perfetto returned: {output}.") + startup_time: float = 2.0 if self.all_heaps != "false" else 0.2 + time.sleep(startup_time) # give it time to spin up + return output + + def _copyPerfDataToHost(self): + self.platform.moveFilesFromPlatform( + os.path.join(self.trace_file_device), + os.path.join(self.host_output_dir), + ) + self.platform.moveFilesFromPlatform( + os.path.join(self.config_file_device), + os.path.join(self.host_output_dir), + ) + + def _generateReport(self): + """Generate an html report from perfetto data.""" + # TODO: implement diff --git a/benchmarking/profilers/perfetto/perfetto_config.py b/benchmarking/profilers/perfetto/perfetto_config.py new file mode 100644 index 00000000..3573633a --- /dev/null +++ b/benchmarking/profilers/perfetto/perfetto_config.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +# Copyright 2004-present Facebook. All Rights Reserved. + + +# duration_ms: {duration_ms} +# max_file_size_bytes: 10000000000 + +CONFIG_TEMPLATE = """ +buffers: {{ + size_kb: {buffer_size_kb} + fill_policy: RING_BUFFER +}} +buffers: {{ + size_kb: {buffer_size2_kb} + fill_policy: RING_BUFFER +}} +{power_config} +{heapprofd_config} +{linux_process_stats_config} +{linux_ftrace_config} +{android_log_config} +{track_event_config} +write_into_file: true +file_write_period_ms: 2500 +max_file_size_bytes: {max_file_size_bytes} +flush_period_ms: 30000 +incremental_state_config {{ + clear_period_ms: 5000 +}} +""" + +POWER_CONFIG = """ +data_sources: {{ + config {{ + name: "android.power" + android_power_config {{ + battery_poll_ms: {battery_poll_ms} + battery_counters: BATTERY_COUNTER_CAPACITY_PERCENT + battery_counters: BATTERY_COUNTER_CHARGE + battery_counters: BATTERY_COUNTER_CURRENT + collect_power_rails: true + }} + }} +}} +""" + +HEAPPROFD_CONFIG = """ +data_sources: {{ + config {{ + name: "android.heapprofd" + target_buffer: 0 + heapprofd_config {{ + sampling_interval_bytes: {sampling_interval_bytes} + process_cmdline: "{app_name}" + shmem_size_bytes: {shmem_size_bytes} + block_client: true + {all_heaps} + }} + }} +}} +""" + +ANDROID_LOG_CONFIG = """ +data_sources: { + config { + name: "android.log" + target_buffer: 0 + android_log_config { + min_prio: PRIO_INFO + log_ids: LID_DEFAULT + log_ids: LID_RADIO + log_ids: LID_EVENTS + log_ids: LID_SYSTEM + log_ids: LID_CRASH + log_ids: LID_KERNEL + } + } +} +""" +LINUX_PROCESS_STATS_CONFIG = """ +data_sources: {{ + config {{ + name: "linux.process_stats" + target_buffer: 0 + process_stats_config {{ + scan_all_processes_on_start: true + proc_stats_poll_ms: 1000 + }} + }} +}} +""" + +LINUX_FTRACE_CONFIG = """ +data_sources: {{ + config {{ + name: "linux.ftrace" + ftrace_config {{ + ftrace_events: "sched/sched_switch" + ftrace_events: "sched/sched_wakeup_new" + ftrace_events: "sched/sched_waking" + ftrace_events: "power/cpu_frequency" + ftrace_events: "power/cpu_idle" + ftrace_events: "power/suspend_resume" + + atrace_apps: "{app_name}" + }} + }} +}} +""" + +TRACK_EVENT_CONFIG = """ +data_sources: {{ + config {{ + name: "track_event" + target_buffer: 1 + }} +}} +""" diff --git a/benchmarking/profilers/utilities.py b/benchmarking/profilers/utilities.py new file mode 100644 index 00000000..d329b46f --- /dev/null +++ b/benchmarking/profilers/utilities.py @@ -0,0 +1,44 @@ +#!/usr/bin/env python + +############################################################################## +# Copyright 2021-present, Facebook, Inc. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +############################################################################## + +import os +from typing import Dict, Mapping +from uuid import uuid4 + +from bridge.file_storage.upload_files.file_uploader import FileUploader +from utils.custom_logger import getLogger + + +def generate_perf_filename(model_name="benchmark", hash=None): + """Given the provided model name and optional hash, generate a unique base filename.""" + if hash is None: + hash = uuid4() + return f"{model_name}_perf_{hash}" + + +def upload_profiling_reports(files: Mapping[str, str]) -> Dict: + """ + Upload to aibench profiling reports manifold bucket. + Accepts dict of key -> local file path, uploads using file basename + and returns meta dict of key -> manifold_url. + """ + meta = {} + profiling_reports_uploader = FileUploader("profiling_reports").get_uploader() + for key, file in files.items(): + if not os.path.isfile(file): + raise FileNotFoundError(f"File {file} does not exist.") + try: + url = profiling_reports_uploader.upload_file(file) + meta.update({key: url}) + except Exception as e: + getLogger().exception( + f"Warning: could not upload {key}: {file}. Skipping.\nException: {e}" + ) + return meta diff --git a/benchmarking/utils/utilities.py b/benchmarking/utils/utilities.py index a79187cc..2387ad7f 100644 --- a/benchmarking/utils/utilities.py +++ b/benchmarking/utils/utilities.py @@ -19,6 +19,7 @@ import sys import tempfile import uuid +import zipfile from time import sleep import certifi @@ -369,3 +370,27 @@ def unpackAdhocFile(configName="generic"): f.write(stream.read()) return path, True + + +def zip_files(input, output: str): + """ + Archive files or folder for uploading. + Input can be file/folder path or list of paths. + Folder hierarchy will be preserved at the folder basename level. + """ + if not isinstance(input, list): + input = [input] + with zipfile.ZipFile(output, "w") as zf: + for path in input: + if os.path.isfile(path): + zf.write(path, os.path.basename(path)) + elif os.path.isdir(path): + for directory, _, files in os.walk(path): + arcdir = directory[directory.find(os.path.basename(path)) :] + zf.write(directory, arcdir) + for f in files: + fpath = os.path.join(directory, f) + arcfpath = os.path.join(arcdir, f) + zf.write(fpath, arcfpath) + else: + raise IOError(f"Could not zip files. {path} is not a valid path.")