GH action to generate report #199
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,35 @@ | ||
| apiVersion: batch/v1 | ||
| kind: Job | ||
| metadata: | ||
| name: disk-usage-report-job | ||
| namespace: jupyterhub | ||
| spec: | ||
| backoffLimit: 0 # No retry on failure | ||
| template: | ||
| metadata: | ||
| labels: | ||
| app: disk-usage-report | ||
| spec: | ||
| containers: | ||
| - name: disk-usage-report | ||
| image: dandiarchive/dandihub-report-generator:latest | ||
| args: | ||
| - "/home/" | ||
| volumeMounts: | ||
| - name: persistent-storage | ||
| mountPath: "/home" | ||
| subPath: "home" | ||
| restartPolicy: Never | ||
| nodeSelector: | ||
| NodeGroupType: default | ||
| NodePool: default | ||
| hub.jupyter.org/node-purpose: user | ||
| tolerations: | ||
| - key: "hub.jupyter.org/dedicated" | ||
| operator: "Equal" | ||
| value: "user" | ||
| effect: "NoSchedule" | ||
| volumes: | ||
| - name: persistent-storage | ||
| persistentVolumeClaim: | ||
| claimName: efs-persist |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,20 @@ | ||
| # manifests/hello-world-pod.yaml | ||
| apiVersion: v1 | ||
| kind: Pod | ||
| metadata: | ||
| name: hello-world-pod | ||
| spec: | ||
| containers: | ||
| - name: hello | ||
| image: busybox | ||
| command: ['sh', '-c', 'echo Hello, World! && sleep 30'] | ||
| nodeSelector: | ||
| NodeGroupType: default | ||
| NodePool: default | ||
| hub.jupyter.org/node-purpose: user | ||
| tolerations: | ||
| - key: "hub.jupyter.org/dedicated" | ||
| operator: "Equal" | ||
| value: "user" | ||
| effect: "NoSchedule" | ||
|
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,180 @@ | ||
| #!/usr/bin/env python3 | ||
|
|
||
| import os | ||
| import csv | ||
| import json | ||
| import sys | ||
| import unittest | ||
| from collections import defaultdict | ||
| from pathlib import Path | ||
| from pprint import pprint | ||
| from typing import Iterable | ||
|
|
||
|
|
||
| def propagate_dir(stats, current_parent, previous_parent): | ||
| assert os.path.isabs(current_parent) == os.path.isabs( | ||
| previous_parent | ||
| ), "current_parent and previous_parent must both be abspath or both be relpath" | ||
| highest_common = os.path.commonpath([current_parent, previous_parent]) | ||
| assert highest_common, "highest_common must either be a target directory or /" | ||
|
|
||
| path_to_propagate = os.path.relpath(previous_parent, highest_common) | ||
| # leaves off last to avoid propagating to the path we are propagating from | ||
| nested_dir_list = path_to_propagate.split(os.sep)[:-1] | ||
| # Add each dir count to all ancestors up to highest common dir | ||
| while nested_dir_list: | ||
| working_dir = os.path.join(highest_common, *nested_dir_list) | ||
| stats[working_dir]["file_count"] += stats[previous_parent]["file_count"] | ||
| stats[working_dir]["total_size"] += stats[previous_parent]["total_size"] | ||
| nested_dir_list.pop() | ||
| previous_parent = working_dir | ||
| stats[highest_common]["file_count"] += stats[previous_parent]["file_count"] | ||
| stats[highest_common]["total_size"] += stats[previous_parent]["total_size"] | ||
|
|
||
|
|
||
| def generate_directory_statistics(data: Iterable[str]): | ||
| # Assumes dirs are listed depth first (files are listed prior to directories) | ||
|
|
||
| stats = defaultdict(lambda: {"total_size": 0, "file_count": 0}) | ||
| previous_parent = "" | ||
| for filepath, size, modified, created, error in data: | ||
| # TODO if error is not None: | ||
| this_parent = os.path.dirname(filepath) | ||
| stats[this_parent]["file_count"] += 1 | ||
| stats[this_parent]["total_size"] += int(size) | ||
|
|
||
| if previous_parent == this_parent: | ||
| continue | ||
| # going deeper | ||
| elif not previous_parent or previous_parent == os.path.dirname(this_parent): | ||
| previous_parent = this_parent | ||
| continue | ||
| else: # previous dir done | ||
| propagate_dir(stats, this_parent, previous_parent) | ||
| previous_parent = this_parent | ||
|
|
||
| # Run a final time with the root directory as this parent | ||
| # During final run, leading dir cannot be empty string, propagate_dir requires | ||
| # both to be abspath or both to be relpath | ||
| leading_dir = previous_parent.split(os.sep)[0] or "/" | ||
| propagate_dir(stats, leading_dir, previous_parent) | ||
| return stats | ||
|
|
||
|
|
||
| def iter_file_metadata(file_path): | ||
| """ | ||
| Reads a tsv and returns an iterable that yields one row of file metadata at | ||
| a time, excluding comments. | ||
| """ | ||
| file_path = Path(file_path) | ||
| with file_path.open(mode="r", newline="", encoding="utf-8") as file: | ||
| reader = csv.reader(file, delimiter="\t") | ||
| for row in reader: | ||
| # Skip empty lines or lines starting with '#' | ||
| if not row or row[0].startswith("#"): | ||
| continue | ||
| yield row | ||
|
|
||
| def update_stats(stats, directory, stat): | ||
| stats["total_size"] += stat["total_size"] | ||
| stats["file_count"] += stat["file_count"] | ||
|
|
||
| # Caches track directories, but not report as a whole | ||
| if stats.get("directories") is not None: | ||
| stats["directories"].append(directory) | ||
|
|
||
| def main(): | ||
| if len(sys.argv) != 2: | ||
| print("Usage: python script.py <input_json_file>") | ||
| sys.exit(1) | ||
|
|
||
| input_tsv_file = sys.argv[1] | ||
| username = input_tsv_file.split("-index.tsv")[0] | ||
|
|
||
| data = iter_file_metadata(input_tsv_file) | ||
| stats = generate_directory_statistics(data) | ||
| cache_types = ["pycache", "user_cache", "yarn_cache", "pip_cache", "nwb_cache"] | ||
| report_stats = { | ||
| "total_size": 0, | ||
| "file_count": 0, | ||
| "caches": { | ||
| cache_type: {"total_size": 0, "file_count": 0, "directories": []} | ||
| for cache_type in cache_types | ||
| } | ||
| } | ||
| # print(f"{directory}: File count: {stat['file_count']}, Total Size: {stat['total_size']}") | ||
| for directory, stat in stats.items(): | ||
| if directory.endswith("__pycache__"): | ||
| update_stats(report_stats["caches"]["pycache"], directory, stat) | ||
| elif directory.endswith(f"{username}/.cache"): | ||
| update_stats(report_stats["caches"]["user_cache"], directory, stat) | ||
| elif directory.endswith(".cache/yarn"): | ||
| update_stats(report_stats["caches"]["yarn_cache"], directory, stat) | ||
| elif directory.endswith(".cache/pip"): | ||
| update_stats(report_stats["caches"]["pip_cache"], directory, stat) | ||
| elif directory == username: | ||
| update_stats(report_stats, username, stat) | ||
|
|
||
| pprint(report_stats) | ||
|
|
||
|
|
||
| class TestDirectoryStatistics(unittest.TestCase): | ||
| def test_propagate_dir(self): | ||
| stats = defaultdict(lambda: {"total_size": 0, "file_count": 0}) | ||
| stats["a/b/c"] = {"total_size": 100, "file_count": 3} | ||
| stats["a/b"] = {"total_size": 10, "file_count": 0} | ||
| stats["a"] = {"total_size": 1, "file_count": 0} | ||
|
|
||
| propagate_dir(stats, "a", "a/b/c") | ||
| self.assertEqual(stats["a"]["file_count"], 3) | ||
| self.assertEqual(stats["a/b"]["file_count"], 3) | ||
| self.assertEqual(stats["a"]["total_size"], 111) | ||
|
|
||
| def test_propagate_dir_abs_path(self): | ||
| stats = defaultdict(lambda: {"total_size": 0, "file_count": 0}) | ||
| stats["/a/b/c"] = {"total_size": 0, "file_count": 3} | ||
| stats["/a/b"] = {"total_size": 0, "file_count": 0} | ||
| stats["/a"] = {"total_size": 0, "file_count": 0} | ||
|
|
||
| propagate_dir(stats, "/a", "/a/b/c") | ||
| self.assertEqual(stats["/a"]["file_count"], 3) | ||
| self.assertEqual(stats["/a/b"]["file_count"], 3) | ||
|
|
||
| def test_propagate_dir_files_in_all(self): | ||
| stats = defaultdict(lambda: {"total_size": 0, "file_count": 0}) | ||
| stats["a/b/c"] = {"total_size": 0, "file_count": 3} | ||
| stats["a/b"] = {"total_size": 0, "file_count": 2} | ||
| stats["a"] = {"total_size": 0, "file_count": 1} | ||
|
|
||
| propagate_dir(stats, "a", "a/b/c") | ||
| self.assertEqual(stats["a"]["file_count"], 6) | ||
| self.assertEqual(stats["a/b"]["file_count"], 5) | ||
|
|
||
| def test_generate_directory_statistics(self): | ||
| sample_data = [ | ||
| ("a/b/file3.txt", 3456, "2024-12-01", "2024-12-02", "OK"), | ||
| ("a/b/c/file1.txt", 1234, "2024-12-01", "2024-12-02", "OK"), | ||
| ("a/b/c/file2.txt", 2345, "2024-12-01", "2024-12-02", "OK"), | ||
| ("a/b/c/d/file4.txt", 4567, "2024-12-01", "2024-12-02", "OK"), | ||
| ("a/e/file3.txt", 5678, "2024-12-01", "2024-12-02", "OK"), | ||
| ("a/e/f/file1.txt", 6789, "2024-12-01", "2024-12-02", "OK"), | ||
| ("a/e/f/file2.txt", 7890, "2024-12-01", "2024-12-02", "OK"), | ||
| ("a/e/f/g/file4.txt", 8901, "2024-12-01", "2024-12-02", "OK"), | ||
| ] | ||
| stats = generate_directory_statistics(sample_data) | ||
| self.assertEqual(stats["a/b/c/d"]["file_count"], 1) | ||
| self.assertEqual(stats["a/b/c"]["file_count"], 3) | ||
| self.assertEqual(stats["a/b"]["file_count"], 4) | ||
| self.assertEqual(stats["a/e/f/g"]["file_count"], 1) | ||
| self.assertEqual(stats["a/e/f"]["file_count"], 3) | ||
| self.assertEqual(stats["a/e"]["file_count"], 4) | ||
| self.assertEqual(stats["a"]["file_count"], 8) | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| if len(sys.argv) > 1 and sys.argv[1] == "test": | ||
| unittest.main( | ||
| argv=sys.argv[:1] | ||
| ) # Run tests if "test" is provided as an argument | ||
| else: | ||
| main() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
| @@ -0,0 +1,63 @@ | ||||||
| #!/usr/bin/env bash | ||||||
|
|
||||||
| set -e | ||||||
|
|
||||||
| # Load environment variables from the file if they are not already set | ||||||
| ENV_FILE=".ec2-session.env" | ||||||
|
There was a problem hiding this comment.
Suggested change
|
||||||
| if [ -f "$ENV_FILE" ]; then | ||||||
| echo "Loading environment variables from $ENV_FILE..." | ||||||
| source "$ENV_FILE" | ||||||
| else | ||||||
| echo "Warning: Environment file $ENV_FILE not found." | ||||||
| fi | ||||||
|
|
||||||
| # Ensure required environment variables are set | ||||||
| if [ -z "$INSTANCE_ID" ]; then | ||||||
| echo "Error: INSTANCE_ID is not set. Cannot proceed with cleanup." | ||||||
| exit 1 | ||||||
| fi | ||||||
|
|
||||||
| if [ -z "$ALLOC_ID" ]; then | ||||||
| echo "Error: ALLOC_ID is not set. Cannot proceed with cleanup." | ||||||
| exit 1 | ||||||
| fi | ||||||
|
|
||||||
| # Check for AWS CLI and credentials | ||||||
| if ! command -v aws &>/dev/null; then | ||||||
| echo "Error: AWS CLI is not installed. Please install it and configure your credentials." | ||||||
| exit 1 | ||||||
| fi | ||||||
|
|
||||||
| if ! aws sts get-caller-identity &>/dev/null; then | ||||||
| echo "Error: Unable to access AWS. Ensure your credentials are configured correctly." | ||||||
| exit 1 | ||||||
| fi | ||||||
|
|
||||||
| # Terminate EC2 instance | ||||||
| echo "Terminating EC2 instance with ID: $INSTANCE_ID..." | ||||||
| if aws ec2 terminate-instances --instance-ids "$INSTANCE_ID" --no-cli-pager; then | ||||||
| echo "Instance termination initiated. Waiting for the instance to terminate..." | ||||||
| if aws ec2 wait instance-terminated --instance-ids "$INSTANCE_ID"; then | ||||||
| echo "Instance $INSTANCE_ID has been successfully terminated." | ||||||
| else | ||||||
| echo "Warning: Instance $INSTANCE_ID may not have terminated correctly." | ||||||
| fi | ||||||
| else | ||||||
| echo "Warning: Failed to terminate instance $INSTANCE_ID. It may already be terminated." | ||||||
| fi | ||||||
|
|
||||||
| # Release Elastic IP | ||||||
| echo "Releasing Elastic IP with Allocation ID: $ALLOC_ID..." | ||||||
| if aws ec2 release-address --allocation-id "$ALLOC_ID"; then | ||||||
| echo "Elastic IP with Allocation ID $ALLOC_ID has been successfully released." | ||||||
| else | ||||||
| echo "Warning: Failed to release Elastic IP with Allocation ID $ALLOC_ID. It may already be released." | ||||||
| fi | ||||||
|
|
||||||
| # Cleanup environment file | ||||||
| if [ -f "$ENV_FILE" ]; then | ||||||
| echo "Removing environment file $ENV_FILE..." | ||||||
| rm -f "$ENV_FILE" | ||||||
| fi | ||||||
|
|
||||||
| echo "Cleanup complete." | ||||||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.