Typing and class import (#54)

* Typing and class import * Ruff fixes --------- Co-authored-by: serriaguirregaray <[email protected]>
danielgafni · Nov 27, 2024 · dee0585 · dee0585
1 parent a590efd
commit dee0585
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/dagster_ray/kuberay/__init__.py b/dagster_ray/kuberay/__init__.py
@@ -1,13 +1,15 @@
 from dagster_ray.kuberay.configs import RayClusterConfig
 from dagster_ray.kuberay.jobs import cleanup_kuberay_clusters, delete_kuberay_clusters
 from dagster_ray.kuberay.ops import cleanup_kuberay_clusters_op, delete_kuberay_clusters_op
+from dagster_ray.kuberay.pipes import PipesKubeRayJobClient
 from dagster_ray.kuberay.resources import KubeRayCluster, RayClusterClientResource
 from dagster_ray.kuberay.schedules import cleanup_kuberay_clusters_daily
 
 __all__ = [
     "KubeRayCluster",
     "RayClusterConfig",
     "RayClusterClientResource",
+    "PipesKubeRayJobClient",
     "cleanup_kuberay_clusters",
     "delete_kuberay_clusters",
     "cleanup_kuberay_clusters_op",

diff --git a/dagster_ray/kuberay/configs.py b/dagster_ray/kuberay/configs.py
@@ -89,8 +89,8 @@ class RayJobConfig(Config):
     entrypoint_memory: float
     entrypoint_num_gpus: int
     suspend: bool = False
-    annotations: Optional[dict[str, str]] = None
-    labels: Optional[dict[str, str]] = None
+    annotations: Optional[Dict[str, str]] = None
+    labels: Optional[Dict[str, str]] = None
     shutdown_after_job_finishes: bool = True
     ttl_seconds_after_finished: int = 60 * 10  # 10 minutes
     active_deadline_seconds: int = 60 * 60 * 24  # 24 hours