Increase ram for some of the goldsky assets (#2152)

opensource-observer · Sep 16, 2024 · b9dab82 · b9dab82
1 parent d0ba5e3
commit b9dab82
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 6 deletions.
diff --git a/warehouse/oso_dagster/factories/goldsky/assets.py b/warehouse/oso_dagster/factories/goldsky/assets.py
@@ -984,7 +984,27 @@ def materialize_asset(
         "opensource.observer/environment": asset_config.environment,
     }
 
-    op_tags: Dict[str, str] = {}
+    op_tags: Dict[str, Any] = {
+        "dagster-k8s/config": {
+            "container_config": {
+                "resources": {
+                    "requests": {"cpu": "2000m", "memory": "3584Mi"},
+                    "limits": {"cpu": "2000m", "memory": "3584Mi"},
+                },
+            },
+            "pod_spec_config": {
+                "node_selector": {"pool_type": "spot",},
+                "tolerations": [
+                    {
+                        "key": "pool_type",
+                        "operator": "Equal",
+                        "value": "spot",
+                        "effect": "PreferNoSchedule",
+                    }
+                ],
+            },
+        }
+    }
 
     if key_prefix:
         group_name = key_prefix if isinstance(key_prefix, str) else "__".join(list(key_prefix))

diff --git a/warehouse/oso_dagster/utils/tags.py b/warehouse/oso_dagster/utils/tags.py
@@ -1,16 +1,16 @@
-from typing import Mapping, Sequence
+import typing as t
 
 
 def add_tags(
-    tags: Mapping[str, str], additional_tags: Mapping[str, str]
-) -> Mapping[str, str]:
+    tags: t.Mapping[str, t.Any], additional_tags: t.Mapping[str, t.Any]
+) -> t.Mapping[str, t.Any]:
     new_tags = dict(tags)
     new_tags.update(additional_tags)
     return new_tags
 
 
 def add_key_prefix_as_tag(
-    tags: Mapping[str, str], key_prefix: Sequence[str] | str | None
+    tags: t.Mapping[str, t.Any], key_prefix: t.Sequence[str] | str | None
 ):
     if key_prefix:
         return add_tags(
@@ -19,5 +19,5 @@ def add_key_prefix_as_tag(
     return add_tags(tags, {})
 
 
-def key_prefix_to_group_name(key_prefix: Sequence[str] | str):
+def key_prefix_to_group_name(key_prefix: t.Sequence[str] | str):
     return key_prefix if isinstance(key_prefix, str) else "__".join(list(key_prefix))