datachain 0.14.2__py3-none-any.whl → 0.39.0__py3-none-any.whl

datachain/studio.py

@@ -1,10 +1,16 @@
 import asyncio
 import os
 import sys
-from typing import TYPE_CHECKING, Optional
+import warnings
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING
+
+import dateparser
+import tabulate
 
 from datachain.config import Config, ConfigLevel
-from datachain.dataset import QUERY_DATASET_PREFIX
+from datachain.data_storage.job import JobStatus
+from datachain.dataset import QUERY_DATASET_PREFIX, parse_dataset_name
 from datachain.error import DataChainError
 from datachain.remote.studio import StudioClient
 from datachain.utils import STUDIO_URL
@@ -16,6 +22,8 @@ POST_LOGIN_MESSAGE = (
     "Once you've logged in, return here "
     "and you'll be ready to start using DataChain with Studio."
 )
+RETRY_MAX_TIMES = 10
+RETRY_SLEEP_SEC = 1
 
 
 def process_jobs_args(args: "Namespace"):
@@ -35,14 +43,28 @@ def process_jobs_args(args: "Namespace"):
             args.workers,
             args.files,
             args.python_version,
+            args.repository,
             args.req,
             args.req_file,
+            args.priority,
+            args.cluster,
+            args.start_time,
+            args.cron,
+            args.no_wait,
+            args.credentials_name,
         )
 
     if args.cmd == "cancel":
         return cancel_job(args.id, args.team)
     if args.cmd == "logs":
         return show_job_logs(args.id, args.team)
+
+    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}'.")
 
 
@@ -60,14 +82,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", {})
@@ -80,11 +112,13 @@ def set_team(args: "Namespace"):
 def login(args: "Namespace"):
     from dvc_studio_client.auth import StudioAuthError, get_access_token
 
+    from datachain.remote.studio import get_studio_env_variable
+
     config = Config().read().get("studio", {})
     name = args.name
     hostname = (
         args.hostname
-        or os.environ.get("DVC_STUDIO_URL")
+        or get_studio_env_variable("URL")
         or config.get("url")
         or STUDIO_URL
     )
@@ -113,6 +147,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
 
 
@@ -141,7 +176,12 @@ def token():
     print(token)
 
 
-def list_datasets(team: Optional[str] = None, name: Optional[str] = None):
+def list_datasets(team: str | None = None, name: str | None = None):
+    def ds_full_name(ds: dict) -> str:
+        return (
+            f"{ds['project']['namespace']['name']}.{ds['project']['name']}.{ds['name']}"
+        )
+
     if name:
         yield from list_dataset_versions(team, name)
         return
@@ -158,18 +198,22 @@ def list_datasets(team: Optional[str] = None, name: Optional[str] = None):
 
     for d in response.data:
         name = d.get("name")
+        full_name = ds_full_name(d)
         if name and name.startswith(QUERY_DATASET_PREFIX):
             continue
 
         for v in d.get("versions", []):
             version = v.get("version")
-            yield (name, version)
+            yield (full_name, version)
 
 
-def list_dataset_versions(team: Optional[str] = None, name: str = ""):
+def list_dataset_versions(team: str | None = None, name: str = ""):
     client = StudioClient(team=team)
 
-    response = client.dataset_info(name)
+    namespace_name, project_name, name = parse_dataset_name(name)
+    if not namespace_name or not project_name:
+        raise DataChainError(f"Missing namespace or project form dataset name {name}")
+    response = client.dataset_info(namespace_name, project_name, name)
 
     if not response.ok:
         raise DataChainError(response.message)
@@ -183,14 +227,18 @@ def list_dataset_versions(team: Optional[str] = None, name: str = ""):
 
 
 def edit_studio_dataset(
-    team_name: Optional[str],
+    team_name: str | None,
     name: str,
-    new_name: Optional[str] = None,
-    description: Optional[str] = None,
-    labels: Optional[list[str]] = None,
+    namespace: str,
+    project: str,
+    new_name: str | None = None,
+    description: str | None = None,
+    attrs: list[str] | None = None,
 ):
     client = StudioClient(team=team_name)
-    response = client.edit_dataset(name, new_name, description, labels)
+    response = client.edit_dataset(
+        name, namespace, project, new_name, description, attrs
+    )
     if not response.ok:
         raise DataChainError(response.message)
 
@@ -198,13 +246,15 @@ def edit_studio_dataset(
 
 
 def remove_studio_dataset(
-    team_name: Optional[str],
+    team_name: str | None,
     name: str,
-    version: Optional[int] = None,
-    force: Optional[bool] = False,
+    namespace: str,
+    project: str,
+    version: str | None = None,
+    force: bool | None = False,
 ):
     client = StudioClient(team=team_name)
-    response = client.rm_dataset(name, version, force)
+    response = client.rm_dataset(name, namespace, project, version, force)
     if not response.ok:
         raise DataChainError(response.message)
 
@@ -222,42 +272,102 @@ def save_config(hostname, token, level=ConfigLevel.GLOBAL):
     return config.config_file()
 
 
+def parse_start_time(start_time_str: str | None) -> str | None:
+    if not start_time_str:
+        return None
+
+    # dateparser#1246: it explores strptime patterns lacking a year, which
+    # triggers a CPython 3.13 DeprecationWarning. Suppress that noise until a
+    # new dateparser release includes the upstream fix.
+    # https://github.com/scrapinghub/dateparser/issues/1246
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            category=DeprecationWarning,
+            module="dateparser\\.utils\\.strptime",
+        )
+        parsed_datetime = dateparser.parse(start_time_str)
+
+    if parsed_datetime is None:
+        raise DataChainError(
+            f"Could not parse datetime string: '{start_time_str}'. "
+            f"Supported formats include: '2024-01-15 14:30:00', 'tomorrow 3pm', "
+            f"'monday 9am', '2024-01-15T14:30:00Z', 'in 2 hours', etc."
+        )
+
+    # Convert to ISO format string
+    return parsed_datetime.isoformat()
+
+
 def show_logs_from_client(client, job_id):
     # Sync usage
     async def _run():
-        async for message in client.tail_job_logs(job_id):
-            if "logs" in message:
-                for log in message["logs"]:
-                    print(log["message"], end="")
-            elif "job" in message:
-                print(f"\n>>>> Job is now in {message['job']['status']} status.")
-
-    asyncio.run(_run())
+        retry_count = 0
+        latest_status = None
+        processed_statuses = set()
+        while True:
+            async for message in client.tail_job_logs(job_id):
+                if "logs" in message:
+                    for log in message["logs"]:
+                        print(log["message"], end="")
+                elif "job" in message:
+                    latest_status = message["job"]["status"]
+                    if latest_status in processed_statuses:
+                        continue
+                    processed_statuses.add(latest_status)
+                    print(f"\n>>>> Job is now in {latest_status} status.")
+
+            try:
+                if retry_count > RETRY_MAX_TIMES or (
+                    latest_status and JobStatus[latest_status].finished()
+                ):
+                    break
+                await asyncio.sleep(RETRY_SLEEP_SEC)
+                retry_count += 1
+            except KeyError:
+                pass
+
+        return latest_status
+
+    final_status = asyncio.run(_run())
 
     response = client.dataset_job_versions(job_id)
     if not response.ok:
         raise DataChainError(response.message)
 
     response_data = response.data
-    if response_data:
+    if response_data and response_data.get("dataset_versions"):
         dataset_versions = response_data.get("dataset_versions", [])
         print("\n\n>>>> Dataset versions created during the job:")
         for version in dataset_versions:
             print(f"    - {version.get('dataset_name')}@v{version.get('version')}")
     else:
-        print("No dataset versions created during the job.")
+        print("\n\nNo dataset versions created during the job.")
+
+    exit_code_by_status = {
+        "FAILED": 1,
+        "CANCELED": 2,
+    }
+    return exit_code_by_status.get(final_status.upper(), 0) if final_status else 0
 
 
 def create_job(
     query_file: str,
-    team_name: Optional[str],
-    env_file: Optional[str] = None,
-    env: Optional[list[str]] = None,
-    workers: Optional[int] = None,
-    files: Optional[list[str]] = None,
-    python_version: Optional[str] = None,
-    req: Optional[list[str]] = None,
-    req_file: Optional[str] = None,
+    team_name: str | None,
+    env_file: str | None = None,
+    env: list[str] | None = None,
+    workers: int | None = None,
+    files: list[str] | None = None,
+    python_version: str | None = None,
+    repository: str | None = None,
+    req: list[str] | None = None,
+    req_file: str | None = None,
+    priority: int | None = None,
+    cluster: str | None = None,
+    start_time: str | None = None,
+    cron: str | None = None,
+    no_wait: bool | None = False,
+    credentials_name: str | None = None,
 ):
     query_type = "PYTHON" if query_file.endswith(".py") else "SHELL"
     with open(query_file) as f:
@@ -276,6 +386,11 @@ def create_job(
     client = StudioClient(team=team_name)
     file_ids = upload_files(client, files) if files else []
 
+    # Parse start_time if provided
+    parsed_start_time = parse_start_time(start_time)
+    if cron and parsed_start_time is None:
+        parsed_start_time = datetime.now(timezone.utc).isoformat()
+
     response = client.create_job(
         query=query,
         query_type=query_type,
@@ -284,7 +399,13 @@ def create_job(
         query_name=os.path.basename(query_file),
         files=file_ids,
         python_version=python_version,
+        repository=repository,
         requirements=requirements,
+        priority=priority,
+        cluster=cluster,
+        start_time=parsed_start_time,
+        cron=cron,
+        credentials_name=credentials_name,
     )
     if not response.ok:
         raise DataChainError(response.message)
@@ -292,12 +413,17 @@ def create_job(
     if not response.data:
         raise DataChainError("Failed to create job")
 
-    job_id = response.data.get("job", {}).get("id")
+    job_id = response.data.get("id")
+
+    if parsed_start_time or cron:
+        print(f"Job {job_id} is scheduled as a task in Studio.")
+        return 0
+
     print(f"Job {job_id} created")
-    print("Open the job in Studio at", response.data.get("job", {}).get("url"))
+    print("Open the job in Studio at", response.data.get("url"))
     print("=" * 40)
 
-    show_logs_from_client(client, job_id)
+    return 0 if no_wait else show_logs_from_client(client, job_id)
 
 
 def upload_files(client: StudioClient, files: list[str]) -> list[str]:
@@ -305,21 +431,19 @@ def upload_files(client: StudioClient, files: list[str]) -> list[str]:
     for file in files:
         file_name = os.path.basename(file)
         with open(file, "rb") as f:
-            file_content = f.read()
-        response = client.upload_file(file_content, file_name)
+            response = client.upload_file(f, file_name)
         if not response.ok:
             raise DataChainError(response.message)
 
         if not response.data:
             raise DataChainError(f"Failed to upload file {file_name}")
 
-        file_id = response.data.get("blob", {}).get("id")
-        if file_id:
+        if file_id := response.data.get("id"):
             file_ids.append(str(file_id))
     return file_ids
 
 
-def cancel_job(job_id: str, team_name: Optional[str]):
+def cancel_job(job_id: str, team_name: str | None):
     token = Config().read().get("studio", {}).get("token")
     if not token:
         raise DataChainError(
@@ -334,7 +458,32 @@ def cancel_job(job_id: str, team_name: Optional[str]):
     print(f"Job {job_id} canceled")
 
 
-def show_job_logs(job_id: str, team_name: Optional[str]):
+def list_jobs(status: str | None, team_name: str | None, limit: int):
+    client = StudioClient(team=team_name)
+    response = client.get_jobs(status, limit)
+    if not response.ok:
+        raise DataChainError(response.message)
+
+    jobs = response.data or []
+    if not jobs:
+        print("No jobs found")
+        return
+
+    rows = [
+        {
+            "ID": job.get("id"),
+            "Name": job.get("name"),
+            "Status": job.get("status"),
+            "Created at": job.get("created_at"),
+            "Created by": job.get("created_by"),
+        }
+        for job in jobs
+    ]
+
+    print(tabulate.tabulate(rows, headers="keys", tablefmt="grid"))
+
+
+def show_job_logs(job_id: str, team_name: str | None):
     token = Config().read().get("studio", {}).get("token")
     if not token:
         raise DataChainError(
@@ -342,4 +491,32 @@ def show_job_logs(job_id: str, team_name: Optional[str]):
         )
 
     client = StudioClient(team=team_name)
-    show_logs_from_client(client, job_id)
+    return show_logs_from_client(client, job_id)
+
+
+def list_clusters(team_name: str | None):
+    client = StudioClient(team=team_name)
+    response = client.get_clusters()
+    if not response.ok:
+        raise DataChainError(response.message)
+
+    clusters = response.data or []
+    if not clusters:
+        print("No clusters found")
+        return
+
+    rows = [
+        {
+            "ID": cluster.get("id"),
+            "Name": cluster.get("name"),
+            "Status": cluster.get("status"),
+            "Cloud Provider": cluster.get("cloud_provider"),
+            "Cloud Credentials": cluster.get("cloud_credentials"),
+            "Is Active": cluster.get("is_active"),
+            "Is Default": cluster.get("default"),
+            "Max Workers": cluster.get("max_workers"),
+        }
+        for cluster in clusters
+    ]
+
+    print(tabulate.tabulate(rows, headers="keys", tablefmt="grid"))

datachain/toolkit/split.py

@@ -1,7 +1,7 @@
 import random
-from typing import Optional
 
 from datachain import C, DataChain
+from datachain.lib.signal_schema import SignalResolvingError
 
 RESOLUTION = 2**31 - 1  # Maximum positive value for a 32-bit signed integer.
 
@@ -9,7 +9,7 @@ RESOLUTION = 2**31 - 1  # Maximum positive value for a 32-bit signed integer.
 def train_test_split(
     dc: DataChain,
     weights: list[float],
-    seed: Optional[int] = None,
+    seed: int | None = None,
 ) -> list[DataChain]:
     """
     Splits a DataChain into multiple subsets based on the provided weights.
@@ -60,7 +60,10 @@ def train_test_split(
         ```
 
     Note:
-        The splits are random but deterministic, based on Dataset `sys__rand` field.
+        Splits reuse the same best-effort shuffle used by `DataChain.shuffle`. Results
+        are typically repeatable, but earlier operations such as `merge`, `union`, or
+        custom SQL that reshuffle rows can change the outcome between runs. Add order by
+        stable keys first when you need strict reproducibility.
     """
     if len(weights) < 2:
         raise ValueError("Weights should have at least two elements")
@@ -69,16 +72,34 @@ def train_test_split(
 
     weights_normalized = [weight / sum(weights) for weight in weights]
 
+    try:
+        dc.signals_schema.resolve("sys.rand")
+    except SignalResolvingError:
+        dc = dc.persist()
+
     rand_col = C("sys.rand")
     if seed is not None:
         uniform_seed = random.Random(seed).randrange(1, RESOLUTION)  # noqa: S311
         rand_col = (rand_col % RESOLUTION) * uniform_seed  # type: ignore[assignment]
     rand_col = rand_col % RESOLUTION  # type: ignore[assignment]
 
-    return [
-        dc.filter(
-            rand_col >= round(sum(weights_normalized[:index]) * (RESOLUTION - 1)),
-            rand_col < round(sum(weights_normalized[: index + 1]) * (RESOLUTION - 1)),
-        )
-        for index, _ in enumerate(weights_normalized)
-    ]
+    boundaries: list[int] = [0]
+    cumulative = 0.0
+    for weight in weights_normalized[:-1]:
+        cumulative += weight
+        boundary = round(cumulative * RESOLUTION)
+        boundaries.append(min(boundary, RESOLUTION))
+    boundaries.append(RESOLUTION)
+
+    splits: list[DataChain] = []
+    last_index = len(weights_normalized) - 1
+    for index in range(len(weights_normalized)):
+        lower = boundaries[index]
+        if index == last_index:
+            condition = rand_col >= lower
+        else:
+            upper = boundaries[index + 1]
+            condition = (rand_col >= lower) & (rand_col < upper)
+        splits.append(dc.filter(condition))
+
+    return splits