datachain 0.19.1__py3-none-any.whl → 0.19.2__py3-none-any.whl

datachain/cli/parser/job.py

@@ -51,6 +51,13 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
         help="Environment variables in KEY=VALUE format",
     )
 
+    studio_run_parser.add_argument(
+        "--cluster-id",
+        type=int,
+        action="store",
+        help="Compute cluster ID to run the job on",
+    )
+
     studio_run_parser.add_argument(
         "--workers",
         type=int,
@@ -165,3 +172,21 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
         default=None,
         help="Team to check logs for (default: from config)",
     )
+
+    studio_clusters_help = "List compute clusters in Studio"
+    studio_clusters_description = "List compute clusters in Studio."
+
+    studio_clusters_parser = jobs_subparser.add_parser(
+        "clusters",
+        parents=[parent_parser],
+        description=studio_clusters_description,
+        help=studio_clusters_help,
+        formatter_class=CustomHelpFormatter,
+    )
+
+    studio_clusters_parser.add_argument(
+        "--team",
+        action="store",
+        default=None,
+        help="Team to list clusters for (default: from config)",
+    )

datachain/cli/parser/studio.py

@@ -89,8 +89,13 @@ def add_auth_parser(subparsers, parent_parser) -> None:
         help="Remove the token from the local project config",
     )
 
-    auth_team_help = "Set default team for Studio operations"
-    auth_team_description = "Set the default team for Studio operations."
+    auth_team_help = "Set or show default team for Studio operations"
+    auth_team_description = (
+        "Set or show the default team for Studio operations. "
+        "This will be used globally by default. "
+        "Use --local to set the team locally for the current project. "
+        "If no team name is provided, the default team will be shown."
+    )
 
     team_parser = auth_subparser.add_parser(
         "team",
@@ -102,13 +107,15 @@ def add_auth_parser(subparsers, parent_parser) -> None:
     team_parser.add_argument(
         "team_name",
         action="store",
+        default=None,
+        nargs="?",
         help="Name of the team to set as default",
     )
     team_parser.add_argument(
-        "--global",
+        "--local",
         action="store_true",
         default=False,
-        help="Set team globally for all projects",
+        help="Set team locally for the current project",
     )
 
     auth_token_help = "View Studio authentication token"  # noqa: S105

datachain/data_storage/schema.py

@@ -13,7 +13,16 @@ from sqlalchemy.sql import func as f
 from sqlalchemy.sql.expression import false, null, true
 
 from datachain.sql.functions import path as pathfunc
-from datachain.sql.types import Int, SQLType, UInt64
+from datachain.sql.types import (
+    JSON,
+    Boolean,
+    DateTime,
+    Int,
+    Int64,
+    SQLType,
+    String,
+    UInt64,
+)
 
 if TYPE_CHECKING:
     from sqlalchemy.engine.interfaces import Dialect
@@ -272,6 +281,19 @@ class DataTable:
             ),
         ]
 
+    @classmethod
+    def listing_columns(cls):
+        return [
+            sa.Column("file__source", String()),
+            sa.Column("file__path", String()),
+            sa.Column("file__size", Int64()),
+            sa.Column("file__version", String()),
+            sa.Column("file__etag", String()),
+            sa.Column("file__is_latest", Boolean()),
+            sa.Column("file__last_modified", DateTime()),
+            sa.Column("file__location", JSON()),
+        ]
+
     def dir_expansion(self):
         return DirExpansion(self.column)
 

datachain/lib/dc/listings.py

@@ -3,19 +3,58 @@ from typing import (
     Optional,
 )
 
+from datachain.lib.listing import LISTING_PREFIX, ls
 from datachain.lib.listing_info import ListingInfo
+from datachain.lib.settings import Settings
+from datachain.lib.signal_schema import SignalSchema
 from datachain.query import Session
+from datachain.query.dataset import DatasetQuery, QueryStep, step_result
 
 from .values import read_values
 
 if TYPE_CHECKING:
     from typing_extensions import ParamSpec
 
+    from datachain.dataset import DatasetVersion
+    from datachain.query.dataset import StepResult
+
     from .datachain import DataChain
 
     P = ParamSpec("P")
 
 
+class ReadOnlyQueryStep(QueryStep):
+    """
+    This step is used to read the dataset in read-only mode.
+    It is used to avoid the need to read the table metadata from the warehouse.
+    This is useful when we want to list the files in the dataset.
+    """
+
+    def apply(self) -> "StepResult":
+        import sqlalchemy as sa
+
+        def q(*columns):
+            return sa.select(*columns)
+
+        table_name = self.catalog.warehouse.dataset_table_name(
+            self.dataset_name, self.dataset_version
+        )
+        dataset_row_cls = self.catalog.warehouse.schema.dataset_row_cls
+        table = dataset_row_cls.new_table(
+            table_name,
+            columns=(
+                [
+                    *dataset_row_cls.sys_columns(),
+                    *dataset_row_cls.listing_columns(),
+                ]
+            ),
+        )
+
+        return step_result(
+            q, table.columns, dependencies=[(self.dataset_name, self.dataset_version)]
+        )
+
+
 def listings(
     session: Optional[Session] = None,
     in_memory: bool = False,
@@ -41,3 +80,75 @@ def listings(
         output={column: ListingInfo},
         **{column: catalog.listings()},  # type: ignore[arg-type]
     )
+
+
+def read_listing_dataset(
+    name: str,
+    version: Optional[str] = None,
+    path: str = "",
+    session: Optional["Session"] = None,
+    settings: Optional[dict] = None,
+) -> tuple["DataChain", "DatasetVersion"]:
+    """Read a listing dataset and return a DataChain and listing version.
+
+    Args:
+        name: Name of the dataset
+        version: Version of the dataset
+        path: Path within the listing to read. Path can have globs.
+        session: Optional Session object to use for reading
+        settings: Optional settings dictionary to use for reading
+
+    Returns:
+        tuple[DataChain, DatasetVersion]: A tuple containing:
+            - DataChain configured for listing files
+            - DatasetVersion object for the specified listing version
+
+    Example:
+        ```py
+        import datachain as dc
+        chain, listing_version = dc.read_listing_dataset(
+            "lst__s3://my-bucket/my-path", version="1.0.0", path="my-path"
+        )
+        chain.show()
+        ```
+    """
+    # Configure and return a DataChain for reading listing dataset files
+    # Uses ReadOnlyQueryStep to avoid warehouse metadata lookups
+    from datachain.lib.dc import Sys
+    from datachain.lib.file import File
+
+    from .datachain import DataChain
+
+    if not name.startswith(LISTING_PREFIX):
+        name = LISTING_PREFIX + name
+
+    session = Session.get(session)
+    dataset = session.catalog.get_dataset(name)
+    if version is None:
+        version = dataset.latest_version
+
+    query = DatasetQuery(
+        name=name,
+        session=session,
+        indexing_column_types=File._datachain_column_types,
+        fallback_to_studio=False,
+    )
+    if settings:
+        cfg = {**settings}
+        if "prefetch" not in cfg:
+            cfg["prefetch"] = 0
+        _settings = Settings(**cfg)
+    else:
+        _settings = Settings(prefetch=0)
+    signal_schema = SignalSchema({"sys": Sys, "file": File})
+
+    query.starting_step = ReadOnlyQueryStep(query.catalog, name, version)
+    query.version = version
+    # We already know that this is a listing dataset,
+    # so we can set the listing function to True
+    query.set_listing_fn(lambda: True)
+
+    chain = DataChain(query, _settings, signal_schema)
+    chain = ls(chain, path, recursive=True, column="file")
+
+    return chain, dataset.get_version(version)

datachain/remote/studio.py

@@ -30,6 +30,7 @@ DatasetExportSignedUrls = Optional[list[str]]
 FileUploadData = Optional[dict[str, Any]]
 JobData = Optional[dict[str, Any]]
 JobListData = dict[str, Any]
+ClusterListData = dict[str, Any]
 logger = logging.getLogger("datachain")
 
 DATASET_ROWS_CHUNK_SIZE = 8192
@@ -391,6 +392,7 @@ class StudioClient:
         requirements: Optional[str] = None,
         repository: Optional[str] = None,
         priority: Optional[int] = None,
+        cluster_id: Optional[int] = None,
     ) -> Response[JobData]:
         data = {
             "query": query,
@@ -403,6 +405,7 @@ class StudioClient:
             "requirements": requirements,
             "repository": repository,
             "priority": priority,
+            "compute_cluster_id": cluster_id,
         }
         return self._send_request("datachain/job", data)
 
@@ -423,3 +426,6 @@ class StudioClient:
     ) -> Response[JobData]:
         url = f"datachain/job/{job_id}/cancel"
         return self._send_request(url, data={}, method="POST")
+
+    def get_clusters(self) -> Response[ClusterListData]:
+        return self._send_request("datachain/clusters", {}, method="GET")

datachain/studio.py

@@ -41,6 +41,7 @@ def process_jobs_args(args: "Namespace"):
             args.req,
             args.req_file,
             args.priority,
+            args.cluster_id,
         )
 
     if args.cmd == "cancel":
@@ -51,6 +52,9 @@ def process_jobs_args(args: "Namespace"):
     if args.cmd == "ls":
         return list_jobs(args.status, args.team, args.limit)
 
+    if args.cmd == "clusters":
+        return list_clusters(args.team)
+
     raise DataChainError(f"Unknown command '{args.cmd}'.")
 
 
@@ -68,14 +72,24 @@ def process_auth_cli_args(args: "Namespace"):
         return logout(args.local)
     if args.cmd == "token":
         return token()
-
     if args.cmd == "team":
         return set_team(args)
     raise DataChainError(f"Unknown command '{args.cmd}'.")
 
 
 def set_team(args: "Namespace"):
-    level = ConfigLevel.GLOBAL if args.__dict__.get("global") else ConfigLevel.LOCAL
+    if args.team_name is None:
+        config = Config().read().get("studio", {})
+        team = config.get("team")
+        if team:
+            print(f"Default team is '{team}'")
+            return 0
+
+        raise DataChainError(
+            "No default team set. Use `datachain auth team <team_name>` to set one."
+        )
+
+    level = ConfigLevel.LOCAL if args.local else ConfigLevel.GLOBAL
     config = Config(level)
     with config.edit() as conf:
         studio_conf = conf.get("studio", {})
@@ -121,6 +135,7 @@ def login(args: "Namespace"):
     level = ConfigLevel.LOCAL if args.local else ConfigLevel.GLOBAL
     config_path = save_config(hostname, access_token, level=level)
     print(f"Authentication complete. Saved token to {config_path}.")
+    print("You can now use 'datachain auth team' to set the default team.")
     return 0
 
 
@@ -268,6 +283,7 @@ def create_job(
     req: Optional[list[str]] = None,
     req_file: Optional[str] = None,
     priority: Optional[int] = None,
+    cluster_id: Optional[int] = None,
 ):
     query_type = "PYTHON" if query_file.endswith(".py") else "SHELL"
     with open(query_file) as f:
@@ -297,6 +313,7 @@ def create_job(
         repository=repository,
         requirements=requirements,
         priority=priority,
+        cluster_id=cluster_id,
     )
     if not response.ok:
         raise DataChainError(response.message)
@@ -380,3 +397,29 @@ def show_job_logs(job_id: str, team_name: Optional[str]):
 
     client = StudioClient(team=team_name)
     show_logs_from_client(client, job_id)
+
+
+def list_clusters(team_name: Optional[str]):
+    client = StudioClient(team=team_name)
+    response = client.get_clusters()
+    if not response.ok:
+        raise DataChainError(response.message)
+
+    clusters = response.data.get("clusters", [])
+    if not clusters:
+        print("No clusters found")
+        return
+
+    rows = [
+        {
+            "ID": cluster.get("id"),
+            "Status": cluster.get("status"),
+            "Cloud Provider": cluster.get("cloud_provider"),
+            "Cloud Credentials": cluster.get("cloud_credentials"),
+            "Is Active": cluster.get("is_active"),
+            "Max Workers": cluster.get("max_workers"),
+        }
+        for cluster in clusters
+    ]
+
+    print(tabulate.tabulate(rows, headers="keys", tablefmt="grid"))

{datachain-0.19.1.dist-info → datachain-0.19.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.19.1
+Version: 0.19.2
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.19.1.dist-info → datachain-0.19.2.dist-info}/RECORD

@@ -15,7 +15,7 @@ datachain/progress.py,sha256=lRzxoYP4Qv2XBwD78sOkmYRzHFpZ2ExVNJF8wAeICtY,770
 datachain/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/script_meta.py,sha256=V-LaFOZG84pD0Zc0NvejYdzwDgzITv6yHvAHggDCnuY,4978
 datachain/semver.py,sha256=UB8GHPBtAP3UJGeiuJoInD7SK-DnB93_Xd1qy_CQ9cU,2074
-datachain/studio.py,sha256=1J2ANFVVA1ysPxBuLibQSnSXt0U9Vfgz9ZNGikYtWdk,11038
+datachain/studio.py,sha256=JCp8qaQSKNs0yiUu2ZGNY7BHlw3CR-OcDjPzqyeMzRU,12369
 datachain/telemetry.py,sha256=0A4IOPPp9VlP5pyW9eBfaTK3YhHGzHl7dQudQjUAx9A,994
 datachain/utils.py,sha256=DNqOi-Ydb7InyWvD9m7_yailxz6-YGpZzh00biQaHNo,15305
 datachain/catalog/__init__.py,sha256=cMZzSz3VoUi-6qXSVaHYN-agxQuAcz2XSqnEPZ55crE,353
@@ -33,8 +33,8 @@ datachain/cli/commands/misc.py,sha256=c0DmkOLwcDI2YhA8ArOuLJk6aGzSMZCiKL_E2JGibV
 datachain/cli/commands/query.py,sha256=Xzfgh14nPVH-sclqX1tpZqgfdTugw5s_44v0D33z6FA,1505
 datachain/cli/commands/show.py,sha256=Cf8wBs12h-xtdOzjU5GTDy2C8rF5HJSF0hDJYER1zH8,1606
 datachain/cli/parser/__init__.py,sha256=sjCIcosAtZqa0m50GMQHqmCkZSYxKyZNwQ29XwRQlP0,15913
-datachain/cli/parser/job.py,sha256=acdVYuTsqluRDI_FYhZ1ohjQcVtBj-taUm8y9tGb0_0,4580
-datachain/cli/parser/studio.py,sha256=Y-1OlQGecLVi9QofvWUfSlPd2ISyaESf7QFGZqGsrdw,3609
+datachain/cli/parser/job.py,sha256=9mEkbhXFIGZxDAjZT9lWWIf3G-HiuJd1oTJNjSrEppc,5295
+datachain/cli/parser/studio.py,sha256=Bo__LKM7qhJGgkyX8M_bCvgZ2Gvqq6r_X4t1NdtaBIY,3881
 datachain/cli/parser/utils.py,sha256=rETdD-9Hq9A4OolgfT7jQw4aoawtbfmkdtH6E7nkhpI,2888
 datachain/client/__init__.py,sha256=1kDpCPoibMXi1gExR4lTLc5pi-k6M5TANiwtXkPoLhU,49
 datachain/client/azure.py,sha256=7yyAgANHfu9Kfh187MKNTT1guvu9Q-WYsi4vYoY3aew,3270
@@ -48,7 +48,7 @@ datachain/data_storage/__init__.py,sha256=9Wit-oe5P46V7CJQTD0BJ5MhOa2Y9h3ddJ4VWT
 datachain/data_storage/db_engine.py,sha256=n8ojCbvVMPY2e3SG8fUaaD0b9GkVfpl_Naa_6EiHfWg,3788
 datachain/data_storage/job.py,sha256=9r0OGwh22bHNIvLHqg8_-eJSP1YYB-BN5HOla5TdCxw,402
 datachain/data_storage/metastore.py,sha256=1PaRTQbL7kjcU1BVjiLjXJLrrLzQtUvpqLmm0pwc1rU,39882
-datachain/data_storage/schema.py,sha256=asZYz1cg_WKfe2Q-k5W51E2z2CzHU5B4QEDZDMFr8yo,9346
+datachain/data_storage/schema.py,sha256=o3JbURKXRg3IJyIVA4QjHHkn6byRuz7avbydU2FlvNY,9897
 datachain/data_storage/serializer.py,sha256=6G2YtOFqqDzJf1KbvZraKGXl2XHZyVml2krunWUum5o,927
 datachain/data_storage/sqlite.py,sha256=BB8x7jtBmHK9lwn2zTo4HgfTKWGF43JxOsGr38J8YV8,25698
 datachain/data_storage/warehouse.py,sha256=imPm4R2V7TkqgGNSO2FGnKu03axU9UVLMfdUPfpwgHE,31747
@@ -103,7 +103,7 @@ datachain/lib/dc/datachain.py,sha256=cQjq6_OWQ_1JKvIqb8snl6mKfuBbpllPEao5ygVINog
 datachain/lib/dc/datasets.py,sha256=g_bBGCUwAwNJypYSUQvrDDqnaw7nfXpvrEvUVPtWATY,11268
 datachain/lib/dc/hf.py,sha256=PJl2wiLjdRsMz0SYbLT-6H8b-D5i2WjeH7li8HHOk_0,2145
 datachain/lib/dc/json.py,sha256=dNijfJ-H92vU3soyR7X1IiDrWhm6yZIGG3bSnZkPdAE,2733
-datachain/lib/dc/listings.py,sha256=2na9v63xO1vPUNaoBSzA-TSN49V7zQAb-4iS1wOPLFE,1029
+datachain/lib/dc/listings.py,sha256=8kX-eELQGHDuOAtavLRJ2iwXkdJQ2bIAv_Z5mKYDJbI,4667
 datachain/lib/dc/pandas.py,sha256=ObueUXDUFKJGu380GmazdG02ARpKAHPhSaymfmOH13E,1489
 datachain/lib/dc/parquet.py,sha256=zYcSgrWwyEDW9UxGUSVdIVsCu15IGEf0xL8KfWQqK94,1782
 datachain/lib/dc/records.py,sha256=J1I69J2gFIBjRTGr2LG-5qn_rTVzRLcr2y3tVDrmHdg,3068
@@ -131,7 +131,7 @@ datachain/query/session.py,sha256=6_ydvPasurmc5tR11dzFj51DpUAo4NxXP9p4ltoTauc,67
 datachain/query/udf.py,sha256=e753bDJzTNjGFQn1WGTvOAWSwjDbrFI1-_DDWkWN2ls,1343
 datachain/query/utils.py,sha256=HaSDNH_XGvp_NIcXjcB7j4vJRPi4_tbztDWclYelHY4,1208
 datachain/remote/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datachain/remote/studio.py,sha256=kUu4TBTKSky1qwdql5DtXNA045qXeoPn4o8G67eZJI4,13666
+datachain/remote/studio.py,sha256=X2uZgPPNDx8mXolp4hajMOqdU_w7jMxRYiki1meMPnw,13919
 datachain/sql/__init__.py,sha256=6SQRdbljO3d2hx3EAVXEZrHQKv5jth0Jh98PogT59No,262
 datachain/sql/selectable.py,sha256=cTc60qVoAwqqss0Vop8Lt5Z-ROnM1XrQmL_GLjRxhXs,1765
 datachain/sql/types.py,sha256=ASSPkmM5EzdRindqj2O7WHLXq8VHAgFYedG8lYfGvVI,14045
@@ -153,9 +153,9 @@ datachain/sql/sqlite/vector.py,sha256=ncW4eu2FlJhrP_CIpsvtkUabZlQdl2D5Lgwy_cbfqR
 datachain/toolkit/__init__.py,sha256=eQ58Q5Yf_Fgv1ZG0IO5dpB4jmP90rk8YxUWmPc1M2Bo,68
 datachain/toolkit/split.py,sha256=ktGWzY4kyzjWyR86dhvzw-Zhl0lVk_LOX3NciTac6qo,2914
 datachain/torch/__init__.py,sha256=gIS74PoEPy4TB3X6vx9nLO0Y3sLJzsA8ckn8pRWihJM,579
-datachain-0.19.1.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.19.1.dist-info/METADATA,sha256=qg4KSU457ARE-A00yjNYNtFP3vhX0yqsxrCGKctXva4,13281
-datachain-0.19.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-datachain-0.19.1.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.19.1.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.19.1.dist-info/RECORD,,
+datachain-0.19.2.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.19.2.dist-info/METADATA,sha256=a4wt0ueB67dUgD1byOAYrqYQZxksLWabVFDYUcv5B9E,13281
+datachain-0.19.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+datachain-0.19.2.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.19.2.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.19.2.dist-info/RECORD,,