Repository: chanzuckerberg/single-cell-curation

Issue:
Support adding labels to larger datasets by reducing memory cost
## Motivation

In order to support larger h5ads we need to reduce the amount of memory required. The way AnnDataLabelAppender is currently written, it triples the amount of memory required. The explosion of memory required can be avoided if we eliminated making a copying of the anndata. 

## Definition of Done

Avoid make a copy of the original dataset in memory to process it with the AnnDataLabelAppender.

### Suggestion

Remove the `copy` and `to_memory` of the annData when initializing AnnDataLabelAppender. The validator which is passed in as a parameter to this class already has the anndata opened and loaded as read only. That means it will be read and modified without need to make a full copy. The required parts will be loaded as needed and written to a new file, leaving the original file intact. 
   
https://github.com/chanzuckerberg/single-cell-curation/blob/79a78b41414441d6f79ecf57baa57e3cea0c5eb3/cellxgene_schema_cli/cellxgene_schema/write_labels.py#L33-L36
changed to 
```
self.adata = validator.adata
```

This change should not affect the contents of the new file produced. 


## Tasks
- [x] remove the copy of anndata into memory
- [x] verify that both the original h5ad remains unchanged.
- [x] verify that the labeled h5ad produced is the same file as before the changes were implemented. No regression.

## Notes
- comparing the hash of the h5ads is enough to ensure they are they same.


Discussion:

Resolving patch:
diff --git a/cellxgene_schema_cli/cellxgene_schema/write_labels.py b/cellxgene_schema_cli/cellxgene_schema/write_labels.py
--- a/cellxgene_schema_cli/cellxgene_schema/write_labels.py
+++ b/cellxgene_schema_cli/cellxgene_schema/write_labels.py
@@ -2,6 +2,7 @@
 import pandas as pd
 
 from typing import List, Dict, Optional
+import traceback
 
 from cellxgene_schema.validate import Validator, ONTOLOGY_CHECKER
 
@@ -30,11 +31,7 @@ def __init__(self, validator: Validator):
                 "Validate AnnData first before attempting to write labels"
             )
 
-        if validator.adata.isbacked:
-            self.adata = validator.adata.to_memory()
-        else:
-            self.adata = validator.adata.copy()
-
+        self.adata = validator.adata
         self.validator = validator
         self.schema_def = validator.schema_def
         self.errors = []
@@ -379,13 +376,13 @@ def write_labels(self, add_labels_file: str):
         try:
             self.adata.write_h5ad(add_labels_file, compression="gzip")
         except Exception as e:
-            self.errors.append(f"Writing h5ad was unsuccessful, got exception '{e}'.")
+            tb = traceback.format_exc()
+            self.errors.append((f"Writing h5ad was unsuccessful, got exception '{e}'.", tb))
 
         # Print errors if any
         if self.errors:
-            self.errors = ["ERROR: " + i for i in self.errors]
-            for e in self.errors:
-                logger.error(e)
+            for e, tb in self.errors:
+                logger.error(e, extra={"exec_info": tb})
             self.was_writing_successful = False
         else:
             self.was_writing_successful = True


Test patch:
diff --git a/cellxgene_schema_cli/tests/test_validate.py b/cellxgene_schema_cli/tests/test_validate.py
--- a/cellxgene_schema_cli/tests/test_validate.py
+++ b/cellxgene_schema_cli/tests/test_validate.py
@@ -1,3 +1,8 @@
+import hashlib
+import os
+import tempfile
+from unittest import mock
+
 import anndata
 import numpy as np
 import unittest
@@ -7,7 +12,7 @@
 from cellxgene_schema.write_labels import AnnDataLabelAppender
 from cellxgene_schema.ontology import OntologyChecker
 from cellxgene_schema.schema import get_schema_definition
-from cellxgene_schema.validate import Validator
+from cellxgene_schema.validate import Validator, validate
 from fixtures.examples_validate import (
     adata_minimal,
     SCHEMA_VERSION,
@@ -17,6 +22,8 @@
     good_var,
     good_uns,
     good_obsm,
+    h5ad_valid,
+    h5ad_invalid,
 )
 
 
@@ -100,9 +107,8 @@ def test_validate_ontology_wrong_term(self):
 
 class TestAddLabelFunctions(unittest.TestCase):
     def setUp(self):
-
         # Set up test data
-        self.test_adata = adata
+        self.test_adata = adata.copy()
         self.test_adata_with_labels = adata_with_labels
         self.schema_def = get_schema_definition(SCHEMA_VERSION)
 
@@ -112,7 +118,6 @@ def setUp(self):
         self.writer = AnnDataLabelAppender(validator)
 
     def test_get_dictionary_mapping_feature_id(self):
-
         # Good
         ids = [
             "ERCC-00002",
@@ -136,7 +141,6 @@ def test_get_dictionary_mapping_feature_id(self):
             self.writer._get_mapping_dict_feature_id(ids)
 
     def test_get_dictionary_mapping_feature_reference(self):
-
         # Good
         ids = [
             "ERCC-00002",
@@ -162,7 +166,6 @@ def test_get_dictionary_mapping_feature_reference(self):
             self.writer._get_mapping_dict_feature_id(ids)
 
     def test_get_dictionary_mapping_curie(self):
-
         # Good
         ids = ["CL:0000066", "CL:0000192"]
         labels = ["epithelial cell", "smooth muscle cell"]
@@ -211,6 +214,21 @@ def test_get_dictionary_mapping_curie(self):
         with self.assertRaises(ValueError):
             self.writer._get_mapping_dict_curie(ids, curie_constraints)
 
+    def test__write__Success(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            labels_path = "/".join([temp_dir, "labels.h5ad"])
+            self.writer.write_labels(labels_path)
+        self.assertTrue(self.writer.was_writing_successful)
+        self.assertFalse(self.writer.errors)
+
+    def test__write__Fail(self):
+        self.writer.adata.write_h5ad = mock.Mock(side_effect=Exception("Test Fail"))
+        with tempfile.TemporaryDirectory() as temp_dir:
+            labels_path = "/".join([temp_dir, "labels.h5ad"])
+            self.writer.write_labels(labels_path)
+        self.assertFalse(self.writer.was_writing_successful)
+        self.assertTrue(self.writer.errors)
+
 
 class TestIgnoreLabelFunctions(unittest.TestCase):
     def setUp(self):
@@ -219,7 +237,6 @@ def setUp(self):
         self.test_adata_with_labels = adata_with_labels
 
     def test_validating_labeled_h5ad_should_fail_if_no_flag_set(self):
-
         validator = Validator()
         validator.adata = self.test_adata_with_labels
         is_valid = validator.validate_adata()
@@ -227,7 +244,6 @@ def test_validating_labeled_h5ad_should_fail_if_no_flag_set(self):
         self.assertFalse(is_valid)
 
     def test_validating_labeled_h5ad_should_pass_if_flag_set(self):
-
         validator = Validator(ignore_labels=True)
         validator.adata = copy.deepcopy(self.test_adata_with_labels)
         is_valid = validator.validate_adata()
@@ -235,6 +251,70 @@ def test_validating_labeled_h5ad_should_pass_if_flag_set(self):
         self.assertTrue(is_valid)
 
 
+class TestValidate(unittest.TestCase):
+    @staticmethod
+    def hash_file(file_name: str) -> str:
+        with open(file_name, "rb") as f:
+            # Read the contents of the file in chunks
+            chunk_size = 1024
+            hasher = hashlib.sha256()
+            while chunk := f.read(chunk_size):
+                hasher.update(chunk)
+        return hasher.hexdigest()
+
+    def test__validate_with_h5ad_valid_and_labels(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            labels_path = "/".join([temp_dir, "labels.h5ad"])
+
+            success, errors, is_seurat_convertible = validate(h5ad_valid, labels_path)
+
+            self.assertTrue(success)
+            self.assertListEqual(errors, [])
+            self.assertTrue(is_seurat_convertible)
+            self.assertTrue(os.path.exists(labels_path))
+            expected_hash = (
+                "55fbc095218a01cad33390f534d6690af0ecd6593f27d7cd4d26e91072ea8835"
+            )
+            actual_hash = self.hash_file(labels_path)
+            original_hash = self.hash_file(h5ad_valid)
+            self.assertNotEqual(
+                original_hash,
+                expected_hash,
+                "Writing labels did not change the dataset from the " "original.",
+            )
+            self.assertEqual(
+                expected_hash,
+                actual_hash,
+                "The shape of the h5ad has changed. Check that the generated file is correct and update the new hash "
+                "to match.",
+            )
+
+    def test__validate_with_h5ad_valid_and_without_labels(self):
+        success, errors, is_seurat_convertible = validate(h5ad_valid)
+
+        self.assertTrue(success)
+        self.assertListEqual(errors, [])
+        self.assertTrue(is_seurat_convertible)
+
+    def test__validate_with_h5ad_invalid_and_with_labels(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            labels_path = "/".join([temp_dir, "labels.h5ad"])
+
+            success, errors, is_seurat_convertible = validate(h5ad_invalid, labels_path)
+
+            self.assertFalse(success)
+            self.assertTrue(errors)
+            self.assertTrue(is_seurat_convertible)
+            self.assertFalse(os.path.exists(labels_path))
+
+    def test__validate_with_h5ad_invalid_and_without_labels(self):
+        success, errors, is_seurat_convertible = validate(h5ad_invalid)
+
+        self.assertFalse(success)
+        self.assertTrue(errors)
+        self.assertTrue(is_seurat_convertible)
+
+
 class TestSeuratConvertibility(unittest.TestCase):
     def validation_helper(self, matrix):
         data = anndata.AnnData(