diff --git a/library/cassandra_hosts.py b/library/cassandra_hosts.py
index db2c2b0..379529b 100644
--- a/library/cassandra_hosts.py
+++ b/library/cassandra_hosts.py
@@ -25,7 +25,7 @@
 def extract_cassandra_groups(inventory_vars, hostvars):
     cassandra_groups = {}
     for name in inventory_vars[GROUPS]:
-        if 'dc-' in name and '-ds' in name:
+        if 'dc-' in name and '-cassandra' in name:
             cassandra_groups[name] = list(inventory_vars[GROUPS][name])
 
     cassandra_ip_mappings= { 'lead_group': '' }
@@ -63,7 +63,7 @@
 
 def prioritize_cassandra_groups(cassandra_groups):
     prioritized_groups = []
-    ds_lead_group = cassandra_groups['lead_group'] + '-ds'
+    ds_lead_group = cassandra_groups['lead_group'] + '-cassandra'
     del cassandra_groups['lead_group']
 
     for ds_ip in cassandra_groups[ds_lead_group]:
diff --git a/tasks/cassandra.yml b/tasks/cassandra.yml
index dfcf458..530db9e 100644
--- a/tasks/cassandra.yml
+++ b/tasks/cassandra.yml
@@ -12,7 +12,7 @@
 - name: Fail if ds group is not provided
   fail:
     msg: 'ds group is not defined in the inventory file, please define the ds group in the inventory file'
-  when: groups['ds'] is not defined
+  when: groups['cassandra'] is not defined
 
 - name: Fail if dc-1 group is not provided
   fail:
@@ -22,7 +22,7 @@
 - name: Fail if dc-1-ds group is not provided
   fail:
     msg: 'dc-1-ds group is not defined in the inventory file, please define the dc-1-ds group in the inventory file'
-  when: groups['dc-1-ds'] is not defined
+  when: groups['dc-1-cassandra'] is not defined
 
 - name: Fail if dc-1-ms group is not provided
   fail: